add original_triples_id to every constructor

besskge/dataset.py

@@ -36,6 +36,11 @@ class KGDataset:
     #: {part: int32[n_triple, {h,r,t}]}
     triples: Dict[str, NDArray[np.int32]]
 
+    #: IDs of the triples in KGDataset.triples wrt
+    #: the ordering in the original array/dataframe
+    #: from where the triples originate.
+    original_triple_ids: Dict[str, NDArray[np.int32]]
+
     #: Entity labels by ID; str[n_entity]
     entity_dict: Optional[List[str]] = None
 
@@ -134,8 +139,8 @@ class should be instantiated manually.
             relation_dict=relation_dict,
             type_offsets=type_offsets,
             triples=triples,
+            original_triple_ids=triple_ids,
         )
-        ds.original_triple_ids = triple_ids  # type: ignore
 
         return ds
 
@@ -228,6 +233,9 @@ def from_dataframe(
                 relation_dict=relation_dict,
                 type_offsets=type_offsets,
                 triples=triples,
+                original_triple_ids={
+                    k: np.arange(v.shape[0]) for k, v in triples.items()
+                },
             )
 
     @classmethod
@@ -289,6 +297,7 @@ def build_ogbl_biokg(cls, root: Path) -> "KGDataset":
             relation_dict=rel_dict,
             type_offsets=type_offsets,
             triples=triples,
+            original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
             neg_heads=neg_heads,
             neg_tails=neg_tails,
         )
@@ -338,6 +347,7 @@ def build_ogbl_wikikg2(cls, root: Path) -> "KGDataset":
             relation_dict=rel_dict,
             type_offsets=None,
             triples=triples,
+            original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
             neg_heads=neg_heads,
             neg_tails=neg_tails,
         )

tests/test_batch_sampler.py

@@ -36,6 +36,7 @@
     relation_dict=None,
     type_offsets=None,
     triples=triples,
+    original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
     neg_heads=None,
     neg_tails=None,
 )

tests/test_bess.py

@@ -79,6 +79,7 @@ def test_bess_inference(
         relation_dict=None,
         type_offsets=None,
         triples=triples,
+        original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
         neg_heads=neg_heads,
         neg_tails=neg_tails,
     )
@@ -306,6 +307,7 @@ def test_bess_topk_prediction(
         relation_dict=None,
         type_offsets=None,
         triples=triples,
+        original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
         neg_heads=neg_heads,
         neg_tails=neg_tails,
     )

tests/test_pipeline.py

@@ -51,6 +51,7 @@ def test_all_scores_pipeline(
         relation_dict=None,
         type_offsets=None,
         triples=triples,
+        original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
     )
 
     partition_mode = "h_shard" if corruption_scheme == "t" else "t_shard"

tests/test_sharding.py

@@ -34,6 +34,7 @@
     relation_dict=None,
     type_offsets={str(i): o for i, o in enumerate(type_offsets)},
     triples=triples,
+    original_triple_ids={k: np.arange(v.shape[0]) for k, v in triples.items()},
     neg_heads=neg_heads,
     neg_tails=neg_tails,
 )