Repository: pykeen/pykeen
Issue: How to get embeddings of entities and relationships by names from a saved trained model?
GitHub labels: question

This issue is derived from #523 . @cthoyt said that it is a common requirement and is better to add to the document.

After reading the previous documents and issues. Here is my solution. Please help to check. Thanks.
```python
import pickle

import torch
from pykeen.datasets.nations import NATIONS_TRAIN_PATH
from pykeen.pipeline import pipeline
from pykeen.triples import TriplesFactory

if __name__ == '__main__':
    # change this path to a TSV file of your dataset
    DATA_PATH = NATIONS_TRAIN_PATH

    # change this path to where your want to store your models
    MODEL_FOLDER = 'my_pykeen_model/'

    tfac = TriplesFactory.from_path(DATA_PATH)
    training, testing, validation = tfac.split([.8, .1, .1])
    result = pipeline(
        training=training,
        testing=testing,
        validation=validation,
        model='TransE',
        stopper='early',
        # short epochs for testing - you should go higher, especially with early stopper enabled
        training_kwargs=dict(num_epochs=5),
    )

    # hibernate trained information and the instance of TriplesFactory
    result.save_to_directory(MODEL_FOLDER)
    with open(MODEL_FOLDER + 'triples_factory.pkl', 'wb') as f:
        pickle.dump(tfac, f)

    # load the trained model and the instance of TriplesFactory
    model = torch.load(MODEL_FOLDER + 'trained_model.pkl')
    with open(MODEL_FOLDER + 'triples_factory.pkl', 'rb') as f:
        loaded_tfac = pickle.load(f)

    # get the embedding of special entities by name
    list_ent_idx = loaded_tfac.entities_to_ids(['china', 'egypt'])
    ent_idx = torch.as_tensor(list_ent_idx, device=model.device)
    entity_embeddings = model.entity_embeddings
    china_embedding, egypt_embedding = entity_embeddings(ent_idx).detach()
    print(china_embedding)
    print(egypt_embedding)

    # get the embedding of special relationships by name
    list_rel_idx = loaded_tfac.relations_to_ids(['embassy', 'officialvisits'])
    rel_idx = torch.as_tensor(list_rel_idx, device=model.device)
    relation_embeddings = model.relation_embeddings
    embassy_embedding, officialvisits_embedding = relation_embeddings(rel_idx).detach()
    print(embassy_embedding)
    print(officialvisits_embedding)

    print('DONE.')
```



The example looks good to me.

For the following explanation, I try to focus on the most relevant parts: using a representation module (aka embedding) to get vectors for given integer IDs, and obtaining these IDs from labels using a TriplesFactory.

### Representation Modules
A representation module maps maps IDs to vectors (or more generally to tensors). Depending on the model, there may be more than one representation, e.g., SimplE has two representations per entity, or TransH having two vectors per relation. You can access entity representations by `model.entity_representations`, which is a list of representation modules. In the simplest case, there is just a single representation module and you can access it via
```python
entity_embedding = model.entity_representations[0]
```
Notice that you need to have integer IDs to use them directly. The representation modules are callable, so we can use
```python
r = entity_embedding(torch.as_tensor([1, 2]))
```
to get the representations for the IDs 1 & 2. The first dimension of the returned tensor will have length 2.

### Obtaining entity/relation IDs from their names
If we start with strings, we first need to convert them to IDs. For this, we can use `TriplesFactory.entity_to_ids` or `TriplesFactory.relation_to_ids`, e.g., 
```python
id_t = torch.as_tensor(triples_factory.entity_to_ids(["china", "egypt"]))
```
It is important to notice that we should use a triples factory with the same mapping that was used to train the model - otherwise we might end up with incorrect IDs.

@mberr , thank you for your detailed and informative descriptions.