cache atom types & bond types once computed

Author: KiddoZhu

torchdrug/data/dataset.py

@@ -414,6 +414,30 @@ class KnowledgeGraphDataset(torch_data.Dataset, core.Configurable):
     The whole dataset contains one knowledge graph.
     """
 
+    def load_triplet(self, triplets, entity_vocab=None, relation_vocab=None, inv_entity_vocab=None,
+                     inv_relation_vocab=None):
+        """
+        Load the dataset from triplets.
+        The mapping between indexes and tokens is specified through either vocabularies or inverse vocabularies.
+
+        Parameters:
+            triplets (array_like): triplets of shape :math:`(n, 3)`
+            entity_vocab (dict of str, optional): maps entity indexes to tokens
+            relation_vocab (dict of str, optional): maps relation indexes to tokens
+            inv_entity_vocab (dict of str, optional): maps tokens to entity indexes
+            inv_relation_vocab (dict of str, optional): maps tokens to relation indexes
+        """
+        entity_vocab, inv_entity_vocab = self._standarize_vocab(entity_vocab, inv_entity_vocab)
+        relation_vocab, inv_relation_vocab = self._standarize_vocab(relation_vocab, inv_relation_vocab)
+
+        num_node = len(entity_vocab) if entity_vocab else None
+        num_relation = len(relation_vocab) if relation_vocab else None
+        self.graph = data.Graph(triplets, num_node=num_node, num_relation=num_relation)
+        self.entity_vocab = entity_vocab
+        self.relation_vocab = relation_vocab
+        self.inv_entity_vocab = inv_entity_vocab
+        self.inv_relation_vocab = inv_relation_vocab
+
     def load_tsv(self, tsv_file, verbose=0):
         """
         Load the dataset from a tsv file.
@@ -483,30 +507,6 @@ def load_tsvs(self, tsv_files, verbose=0):
         self.load_triplet(triplets, inv_entity_vocab=inv_entity_vocab, inv_relation_vocab=inv_relation_vocab)
         self.num_samples = num_samples
 
-    def load_triplet(self, triplets, entity_vocab=None, relation_vocab=None, inv_entity_vocab=None,
-                     inv_relation_vocab=None):
-        """
-        Load the dataset from triplets.
-        The mapping between indexes and tokens is specified through either vocabularies or inverse vocabularies.
-
-        Parameters:
-            triplets (array_like): triplets of shape :math:`(n, 3)`
-            entity_vocab (dict of str, optional): maps entity indexes to tokens
-            relation_vocab (dict of str, optional): maps relation indexes to tokens
-            inv_entity_vocab (dict of str, optional): maps tokens to entity indexes
-            inv_relation_vocab (dict of str, optional): maps tokens to relation indexes
-        """
-        entity_vocab, inv_entity_vocab = self._standarize_vocab(entity_vocab, inv_entity_vocab)
-        relation_vocab, inv_relation_vocab = self._standarize_vocab(relation_vocab, inv_relation_vocab)
-
-        num_node = len(entity_vocab) if entity_vocab else None
-        num_relation = len(relation_vocab) if relation_vocab else None
-        self.graph = data.Graph(triplets, num_node=num_node, num_relation=num_relation)
-        self.entity_vocab = entity_vocab
-        self.relation_vocab = relation_vocab
-        self.inv_entity_vocab = inv_entity_vocab
-        self.inv_relation_vocab = inv_relation_vocab
-
     def _standarize_vocab(self, vocab, inverse_vocab):
         if vocab is not None:
             if isinstance(vocab, dict):
@@ -609,7 +609,7 @@ def round_to_boundary(i):
 
     if key_lengths is not None:
         assert lengths is None
-        key2count = torch.bincount(keys)
+        key2count = keys.bincount()
         key_offset = 0
         lengths = []
         for key_length in key_lengths: