Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Track original triple IDs in KGDataset.from_triples #37

Merged
merged 3 commits into from
Jan 8, 2024
Merged

Track original triple IDs in KGDataset.from_triples #37

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ae821e4
track original triple IDs when shuffling/splitting in KGDataset.from_…
AlCatt91 Jan 3, 2024
6514372
lint fix
AlCatt91 Jan 3, 2024
74ffbb6
add original_triples_id to every constructor
AlCatt91 Jan 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions besskge/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,9 @@ def from_triples(
and relations have already been assigned. Note that, if entities have
types, entities of the same type need to have contiguous IDs.
Triples are randomly split in train/validation/test sets.
The attribute `KGDataset.original_triple_ids` stores the IDs
of the triples in each split wrt the original ordering in `data`.

If a pre-defined train/validation/test split is wanted, the KGDataset
class should be instantiated manually.

Expand All @@ -114,21 +117,27 @@ class should be instantiated manually.
num_valid = int(num_triples * split[1])

rng = np.random.default_rng(seed=seed)
rng.shuffle(data, axis=0)

triples = dict()
triples["train"], triples["valid"], triples["test"] = np.split(
data, (num_train, num_train + num_valid), axis=0
id_shuffle = rng.permutation(np.arange(num_triples))
triple_ids = dict()
triple_ids["train"], triple_ids["valid"], triple_ids["test"] = np.split(
id_shuffle, (num_train, num_train + num_valid), axis=0
)
triples = dict()
triples["train"] = data[triple_ids["train"]]
triples["valid"] = data[triple_ids["valid"]]
triples["test"] = data[triple_ids["test"]]

return cls(
ds = cls(
n_entity=data[:, [0, 2]].max() + 1,
n_relation_type=data[:, 1].max() + 1,
entity_dict=entity_dict,
relation_dict=relation_dict,
type_offsets=type_offsets,
triples=triples,
)
ds.original_triple_ids = triple_ids # type: ignore
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might be more consistent for downstream analytics if original_triple_ids always was a member of the dataset. I'd add a dict with values arange(...) if the dataset is not created through from_triples

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree, good point, thanks! I should have added this to every constructor


return ds

@classmethod
def from_dataframe(
Expand Down