Pre-trained Models

To facilitate the usage of knowledge graph representations in semantic tasks, we provide a bunch of pre-trained embeddings for some common datasets.

Wikidata5m

Wikidata5m is a large-scale knowledge graph dataset constructed from Wikidata and Wikipedia. It contains plenty of entities in the general domain, such as celebrities, events, concepts and things.

We trained 5 standard knowledge graph embedding models on Wikidata5m. The performance benchmark of these models can be found here.

Model	Dimension	Size	Download link
TransE	512	9.33 GB	transe_wikidata5m.pkl
DistMult	512	9.33 GB	distmult_wikidata5m.pkl
ComplEx	512	9.33 GB	complex_wikidata5m.pkl
SimplE	512	9.33 GB	simple_wikidata5m.pkl
RotatE	512	9.33 GB	rotate_wikidata5m.pkl

Load pre-trained models

The pre-trained models can be loaded through pickle.

import pickle
with open("transe_wikidata5m.pkl", "rb") as fin:
    model = pickle.load(fin)
entity2id = model.graph.entity2id
relation2id = model.graph.relation2id
entity_embeddings = model.solver.entity_embeddings
relation_embeddings = model.solver.relation_embeddings

Load the alias mapping from the dataset. Now we can access the embeddings by natural language index.

import graphvite as gv
alias2entity = gv.dataset.wikidata5m.alias2entity
alias2relation = gv.dataset.wikidata5m.alias2relation
print(entity_embeddings[entity2id[alias2entity["machine learning"]]])
print(relation_embeddings[relation2id[alias2relation["field of work"]]])