Initialize the project.

Author: lsvih

.gitignore

@@ -127,3 +127,4 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+.idea

README.md

@@ -1 +1,28 @@
-# AtTGen
+# AtTGen
+
+## Usage
+
+```
+usage: python3 main.py [-h] [--name NAME] [--do_train] [--do_eval]
+               [--data_dir DATA_DIR] [--seed SEED] [--gpu_ids GPU_IDS]
+               [--batch_size BATCH_SIZE] [--lr LR] [--epoch EPOCH]
+               [--emb_dim EMB_DIM] [--encode_dim ENCODE_DIM]
+
+configuration
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --name NAME           Experiment name, for logging and saving models
+  --do_train            Whether to run traininog.
+  --do_eval             Whether to run eval on the test set.
+  --data_dir DATA_DIR   The input data dir.
+  --seed SEED           The random seed for initialization
+  --gpu_ids GPU_IDS     The GPU ids
+  --batch_size BATCH_SIZE
+                        Total batch size for training.
+  --lr LR               The initial learning rate for Adam.
+  --epoch EPOCH         Total number of training epochs to perform.
+  --emb_dim EMB_DIM     The dimension of the embedding
+  --encode_dim ENCODE_DIM
+                        The dimension of the encoding
+```

args.py

@@ -0,0 +1,84 @@
+import argparse
+
+import torch
+
+from utils import get_device
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='configuration')
+
+    parser.add_argument("--name",
+                        default="1",
+                        type=str,
+                        help="Experiment name, for logging and saving models")
+    parser.add_argument("--do_train",
+                        action='store_true',
+                        default=True,
+                        help="Whether to run traininog.")
+    parser.add_argument("--do_eval",
+                        action='store_true',
+                        default=True,
+                        help="Whether to run eval on the test set.")
+    parser.add_argument("--data_dir",
+                        default="./data/CNShipNet",
+                        type=str,
+                        help="The input data dir.")
+    parser.add_argument("--word_vocab",
+                        default="word_vocab.json",
+                        type=str,
+                        help="The vocabulary file.")
+    parser.add_argument("--ontology_vocab",
+                        default="attribute_vocab.json",
+                        type=str,
+                        help="The ontology class file.")
+    parser.add_argument('--seed',
+                        type=int,
+                        default=42,
+                        help="The random seed for initialization")
+    parser.add_argument('--gpu_ids',
+                        type=str,
+                        default='0',
+                        help="The GPU ids")
+
+    # Hyperparameters
+    # Batch size
+    parser.add_argument("--batch_size",
+                        default=512,
+                        type=int,
+                        help="Total batch size for training.")
+    # Learning rate
+    parser.add_argument("--lr",
+                        default=2e-4,
+                        type=float,
+                        help="The initial learning rate for Adam.")
+    # Epochs
+    parser.add_argument("--epoch",
+                        default=40,
+                        type=int,
+                        help="Total number of training epochs to perform.")
+    # emb_dim
+    parser.add_argument("--emb_dim",
+                        default=200,
+                        type=int,
+                        help="The dimension of the embedding")
+    # encode_dim
+    parser.add_argument("--encode_dim",
+                        default=200,
+                        type=int,
+                        help="The dimension of the encoding")
+
+    args = parser.parse_args()
+    if args.gpu_ids == "":
+        n_gpu = 0
+        device = torch.device('cpu')
+    else:
+        gpu_ids = [int(device_id) for device_id in args.gpu_ids.split()]
+        args.gpu_ids = gpu_ids
+        device, n_gpu = get_device(gpu_ids[0])
+        if n_gpu > 1:
+            n_gpu = len(gpu_ids)
+    args.device = device
+    args.n_gpu = n_gpu
+
+    return args

data/.gitkeep

@@ -0,0 +1,255 @@
+import json
+import os
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+from utils import find_entity_id_from_tokens, seq_padding, sort_all
+
+
+class TreeDataset(Dataset):
+    def __init__(self, data_dir: str = './data/CNShipNet', data_type: str = "train",
+                 word_vocab: str = 'word_vocab.json', ontology_vocab: str = 'attribute_vocab.json',
+                 order: List[str] = ("subject", "object", "predicate")):
+        print('Loading {} data...'.format(data_type))
+        self.data_dir = data_dir
+        self.order = order
+        self.word_vocab = json.load(open(os.path.join(data_dir, word_vocab)))
+        self.ontology_vocab = json.load(open(os.path.join(data_dir, ontology_vocab)))
+        vocab_size = len(self.word_vocab)
+        rel_token = {k: (v + vocab_size) for k, v in self.ontology_vocab.items()}
+        self.word_vocab.update(rel_token)
+        self.word_vocab['[pre]'] = len(self.word_vocab)
+        self.ontology_class = list(rel_token.keys())
+
+        self.tokenizer = lambda text: text.split(" ")
+
+        self.text = []
+        self.text_length = []
+        self.spo_list = []
+        self.token_ids = []
+        self.S1 = []
+        self.S2 = []
+        self.S_K1_in = []
+        self.O_K1_in = []
+        self.S_K2_in = []
+        self.O_K2_in = []
+        self.O1 = []
+        self.O2 = []
+        self.P1 = []
+        self.P2 = []
+        self.P_K1_in = []
+        self.P_K2_in = []
+        file = open(os.path.join(self.data_dir, "new_{}_data.json".format(data_type))).read().strip().split('\n')
+        for line in tqdm(file):
+            instance = json.loads(line)
+            if data_type == 'train':
+                expanded_instances = self.spo_to_seq(instance["text"], instance["spo_list"], self.tokenizer,
+                                                     self.ontology_class)
+                instances = expanded_instances
+            else:
+                token = self.tokenizer(instance["text"]) + ['[pre]'] + self.ontology_class
+                instance['text'] = ' '.join(token)
+                instances = [instance]
+            for instance in instances:
+                text = instance["text"]
+                spo_list = instance["spo_list"]
+                text_id = []
+                for c in (self.tokenizer(text)):
+                    text_id.append(self.word_vocab.get(c, self.word_vocab["<oov>"]))
+                self.text_length.append(len(text_id))
+                assert len(text_id) > 0
+                self.token_ids.append(text_id)
+
+                s_k1 = instance.get("s_k1", 0)
+                s_k2 = instance.get("s_k2", 0)
+                o_k1 = instance.get("o_k1", 0)
+                o_k2 = instance.get("o_k2", 0)
+                p_k1 = instance.get("p_k1", 0)
+                p_k2 = instance.get("p_k2", 0)
+
+                s1_gt = instance.get("s1_gt", [])
+                s2_gt = instance.get("s2_gt", [])
+                o1_gt = instance.get("o1_gt", [])
+                o2_gt = instance.get("o2_gt", [])
+                p1_gt = instance.get("p1_gt", [])
+                p2_gt = instance.get("p2_gt", [])
+
+                self.text.append(self.tokenizer(text))
+                self.spo_list.append(spo_list)
+
+                self.S1.append(s1_gt)
+                self.S2.append(s2_gt)
+                self.O1.append(o1_gt)
+                self.O2.append(o2_gt)
+                self.P1.append(p1_gt)
+                self.P2.append(p2_gt)
+                self.S_K1_in.append([s_k1])
+                self.S_K2_in.append([s_k2])
+                self.O_K1_in.append([o_k1])
+                self.O_K2_in.append([o_k2])
+                self.P_K1_in.append([p_k1])
+                self.P_K2_in.append([p_k2])
+
+        self.token_ids = np.array(seq_padding(self.token_ids))
+
+        # training
+        self.S1 = np.array(seq_padding(self.S1))
+        self.S2 = np.array(seq_padding(self.S2))
+        self.O1 = np.array(seq_padding(self.O1))
+        self.O2 = np.array(seq_padding(self.O2))
+        self.P1 = np.array(seq_padding(self.P1))
+        self.P2 = np.array(seq_padding(self.P2))
+
+        # self.K1_in, self.K2_in = np.array(self.K1_in), np.array(self.K2_in)
+        # only two time step are used for training
+        self.S_K1_in = np.array(self.S_K1_in)
+        self.S_K2_in = np.array(self.S_K2_in)
+        self.O_K1_in = np.array(self.O_K1_in)
+        self.O_K2_in = np.array(self.O_K2_in)
+        self.P_K1_in = np.array(self.P_K1_in)
+        self.P_K2_in = np.array(self.P_K2_in)
+
+    def __getitem__(self, index):
+        return (
+            self.token_ids[index],
+            self.S1[index],
+            self.S2[index],
+            self.O1[index],
+            self.O2[index],
+            self.P1[index],
+            self.P2[index],
+            self.S_K1_in[index],
+            self.S_K2_in[index],
+            self.O_K1_in[index],
+            self.O_K2_in[index],
+            self.P_K1_in[index],
+            self.P_K2_in[index],
+            self.text[index],  # original text
+            self.text_length[index],  # token length
+            self.spo_list[index],  # spo list
+        )
+
+    def __len__(self):
+        return len(self.text)
+
+    def spo_to_seq(self, text, spo_list, tokenizer, ontology_class):
+        tree = self.spo_to_tree(spo_list, self.order)
+        tokens = tokenizer(text) + ['[pre]'] + ontology_class
+
+        def to_ent(outp):
+            # side effect!
+            ent1, ent2 = [[0] * len(tokens) for _ in range(2)]
+            for name in outp:
+                _id = find_entity_id_from_tokens(tokens, self.tokenizer(name))
+                ent1[_id] = 1
+                ent2[_id + len(self.tokenizer(name)) - 1] = 1
+            return ent1, ent2
+
+        def to_in_key(inp, name):
+            # side effect!
+            if not inp:
+                return 0, 0
+
+            k1 = find_entity_id_from_tokens(tokens, self.tokenizer(inp))
+            k2 = k1 + len(self.tokenizer(inp)) - 1
+            out = k1, k2
+            return out
+
+        results = []
+        for t in tree:
+            t1_in, t2_in, t1_out, t2_out, t3_out = t
+            for name, ori_out, ori_in in zip(
+                    self.order, (t1_out, t2_out, t3_out), (t1_in, t2_in, None)
+            ):
+                new_out = to_ent(ori_out)
+                if name == "predicate":
+                    p1, p2 = new_out
+                    p_k1, p_k2 = to_in_key(ori_in, name)
+                elif name == "subject":
+                    s1, s2 = new_out
+                    s_k1, s_k2 = to_in_key(ori_in, name)
+                elif name == "object":
+                    o1, o2 = new_out
+                    o_k1, o_k2 = to_in_key(ori_in, name)
+                else:
+                    raise ValueError("should be in predicate, subject, object")
+
+            result = {
+                "text": ' '.join(tokens),
+                "spo_list": spo_list,
+                "s_k1": s_k1,
+                "s_k2": s_k2,
+                "o_k1": o_k1,
+                "o_k2": o_k2,
+                "p_k1": p_k1,
+                "p_k2": p_k2,
+                "s1_gt": s1,
+                "s2_gt": s2,
+                "o1_gt": o1,
+                "o2_gt": o2,
+                "p1_gt": p1,
+                "p2_gt": p2,
+            }
+
+            results.append(result)
+        return results
+
+    def spo_to_tree(self, spo_list: List[Dict[str, str]], order=("subject", "object", "predicate")):
+        """return the ground truth of the tree: rel, subj, obj, used for teacher forcing.
+        r: given text, one of the relations
+        s: given r_1, one of the subjects
+        rel: multi-label classification of relation
+        subj: multi-label classification of subject
+        obj: multi-label classification of object
+        Arguments:
+            spo_list {List[Dict[str, str]]} -- [description]
+        Returns:
+            List[Tuple[str]] -- [(r, s, rel, subj, obj)]
+        """
+        result = []
+        t1_out = list(set(t[order[0]] for t in spo_list))
+        for t1_in in t1_out:
+            t2_out = list(set(t[order[1]] for t in spo_list if t[order[0]] == t1_in))
+            for t2_in in t2_out:
+                t3_out = list(
+                    set(
+                        t[order[2]]
+                        for t in spo_list
+                        if t[order[0]] == t1_in and t[order[1]] == t2_in
+                    )
+                )
+                result.append((t1_in, t2_in, t1_out, t2_out, t3_out))
+        return result
+
+
+def collate(batch: List[Tuple]):
+    batch_data = list(zip(*batch))
+    token_len = batch_data[-2]
+    batch_data, orig_idx = sort_all(batch_data, token_len)
+    token_ids, s1, s2, o1, o2, p1, p2, s_k1_in, s_k2_in, o_k1_in, o_k2_in, p_k1_in, p_k2_in, text, token_len, spo_list = batch_data
+    token_ids = torch.LongTensor(np.array(token_ids))
+    s1 = torch.FloatTensor(np.array(s1))
+    s2 = torch.FloatTensor(np.array(s2))
+    o1 = torch.FloatTensor(np.array(o1))
+    o2 = torch.FloatTensor(np.array(o2))
+    p1 = torch.FloatTensor(np.array(p1))
+    p2 = torch.FloatTensor(np.array(p2))
+    s_k1_in = torch.LongTensor(np.array(s_k1_in))
+    s_k2_in = torch.LongTensor(np.array(s_k2_in))
+    o_k1_in = torch.LongTensor(np.array(o_k1_in))
+    o_k2_in = torch.LongTensor(np.array(o_k2_in))
+    p_k1_in = torch.LongTensor(np.array(p_k1_in))
+    p_k2_in = torch.LongTensor(np.array(p_k2_in))
+    token_len = torch.LongTensor(token_len)
+    return {'token_ids': token_ids, 's1': s1, 's2': s2, 'o1': o1, 'o2': o2, 'p1': p1, 'p2': p2, 's_k1_in': s_k1_in,
+            's_k2_in': s_k2_in, 'o_k1_in': o_k1_in, 'o_k2_in': o_k2_in, 'p_k1_in': p_k1_in, 'p_k2_in': p_k2_in,
+            'text': text, 'token_len': token_len, 'spo_list': spo_list}
+
+
+if __name__ == '__main__':
+    x = TreeDataset()
+    exit()