Repository: pykeen/pykeen
Issue: How to get embeddings of specific entity?
Root cause: Documentation Deficiency

I use "nations" dataset to complete training and get "trained_model.pkl" model. How can I load the local "trained_model.pkl" file and get embeddings of specific entity?
I refer to the answer and complete the code below, but it doesn't work. 

```python
# model = pipeline_result.model   (Maybe it works, but how to load the pipeline_result?)
model = torch.load('./trained_model/nations_TransE_200d/trained_model.pkl')
entity_to_id = model.triples_factory.entity_to_id

entities = ['egypt', 'usa']   ## I want to get entity_embeddings of entity: 'egypt'.
entity_ids = torch.tensor([entity_to_id[entity] for entity in entities], dtype=torch.long)
embedded_entities = entity_embeddings[entity_ids]
```

## Related

- https://github.com/pykeen/pykeen/issues/89


Hi @fiego, this is a similar question to #89. We improved the first steps tutorial as a solution in #336, but since we have been a bit slow to make a new PyKEEN official release, it still lives in the (non-default) "latest" docs instead of the "stable" one. Please have a look at https://pykeen.readthedocs.io/en/latest/tutorial/first_steps.html#using-learned-embeddings and see if that helps. Keep in mind that each model might have a different name/number of embeddings for each entity - the tutorial points out how to avoid a few common issues

If this isn't helpful enough, please let me know and I will do my best to explain further and update the docs accordingly :)

> If this isn't helpful enough, please let me know and I will do my best to explain further and update the docs accordingly :)

Yes, I have referenced this anser, but I am confused, how I can load the pipeline_result from local file.

model = pipeline_result.model (Maybe it works, but how to load the pipeline_result?)

It's not directly possible to load an instance of a PipelineResult from a file. The `PipelineResult.save` function just outputs the stuff you would typically want to re-use. I guess it would be possible to pickle the whole PipelineResult class after running the pipeline if you want access to all of its parts later.

However, since you just want to load up the model, using torch.load is the right thing. The "first steps" tutorial also has a small section on how to do that, but I'm not exactly sure what's going wrong for you. Is an exception being thrown? If so, it would be helpful for you to copy the traceback for me to take a look

Here's a full example using the pipeline to get the Egypt embedding after (a rather fast/poor) training

```python
from pykeen.pipeline import pipeline
import torch

pipeline_result = pipeline(
    dataset='Nations',
    model='PairRE',
    training_kwargs=dict(num_epochs=1, use_tqdm=False),  # make this higher for real usage
)

# This function only operates on lists
list_idx = pipeline_result.training.entities_to_ids(['egypt'])
# [4]

idx = torch.as_tensor(list_idx, device=pipeline_result.model.device)
# tensor([4])

# new-style models like PairRE always have a list of entity representations.
#   just take the first one
entity_embeddings = pipeline_result.model.entity_representations[0]

# entity_embeddings is a pykeen.nn.Embedding, which is a torch.nn.Module
# whose forward function takes a list of entity indices, and returns
# their embeddings. Use detach() to take it out of the gradient mode,
# then slice with [0] to get the first (and only) element
egypt_embedding = entity_embeddings(idx).detach()[0]
```

Here's the same thing with TransE, note that the way you get the embeddings is different. That's just a model-specific thing. Maybe we can put a "default" embedding in each model for the entities to make this more consistent in a future PR

```python
from pykeen.pipeline import pipeline
import torch

pipeline_result = pipeline(
    dataset='Nations',
    model='TransE',
    training_kwargs=dict(num_epochs=1, use_tqdm=False),
)
list_idx = pipeline_result.training.entities_to_ids(['egypt'])
idx = torch.as_tensor(list_idx, device=pipeline_result.model.device)

# This is different!
entity_embeddings = pipeline_result.model.entity_embeddings

egypt_embedding = entity_embeddings(idx).detach()[0]
```



Here's the version where you save the model in the middle. Note that you have to reload the dataset. With newer versions of pykeen, we stopped storing the training triples in the model - this saves a lot of GPU memory and we figured that most people would still have access to the data after training if they needed it (like in your situation where you want to use entity labels instead of identifiers)

```python
from pykeen.pipeline import pipeline
import torch
import pystow

directory = pystow.join('pykeen', 'tmp')
pipeline_result = pipeline(
    dataset='Nations',
    model='TransE',
    training_kwargs=dict(num_epochs=1, use_tqdm=False),
)
pipeline_result.save_to_directory(directory)

# just so we're not tempted to reuse this in the rest of the example
del pipeline_result

model = torch.load(directory / 'trained_model.pkl')

# newer versions of PyKEEN don't save the triples factory in the model
#   anymore, so we need to reload it
from pykeen.datasets import Nations
dataset = Nations()

# now, just access the training data from the dataset and do the rest the same
list_idx = dataset.training.entities_to_ids(['egypt'])
idx = torch.as_tensor(list_idx, device=model.device)
entity_embeddings = model.entity_embeddings
egypt_embedding = entity_embeddings(idx).detach()[0]
egypt_embedding
```

> If this isn't helpful enough, please let me know and I will do my best to explain further and update the docs accordingly :)

Yes, I have referenced this anser, but I am confused, 

model = pipeline_result.model (Maybe it works, but how to load the pipeline_result?)

> Here's the version where you save the model in the middle. Note that you have to reload the dataset. With newer versions of pykeen, we stopped storing the training triples in the model - this saves a lot of GPU memory and we figured that most people would still have access to the data after training if they needed it (like in your situation where you want to use entity labels instead of identifiers)
> 
> ```python
> from pykeen.pipeline import pipeline
> import torch
> import pystow
> 
> directory = pystow.join('pykeen', 'tmp')
> pipeline_result = pipeline(
>     dataset='Nations',
>     model='TransE',
>     training_kwargs=dict(num_epochs=1, use_tqdm=False),
> )
> pipeline_result.save_to_directory(directory)
> 
> # just so we're not tempted to reuse this in the rest of the example
> del pipeline_result
> 
> model = torch.load(directory / 'trained_model.pkl')
> 
> # newer versions of PyKEEN don't save the triples factory in the model
> #   anymore, so we need to reload it
> from pykeen.datasets import Nations
> dataset = Nations()
> 
> # now, just access the training data from the dataset and do the rest the same
> list_idx = dataset.training.entities_to_ids(['egypt'])
> idx = torch.as_tensor(list_idx, device=model.device)
> entity_embeddings = model.entity_embeddings
> egypt_embedding = entity_embeddings(idx).detach()[0]
> egypt_embedding
> ```

Thank you very much,! 
It works, the following is my code after referring to your reply.

```python
### nations
TRAIN_PATH = "./nations/train.txt"
TEST_PATH = "./nations/test.txt"
VALIDATE_PATH = "./nations/valid.txt"

### TriplesFactory: training, testing, validation
training = TriplesFactory.from_path(
    path=TRAIN_PATH,
    create_inverse_triples=False,)

testing = TriplesFactory.from_path(
    path=TEST_PATH,
    create_inverse_triples=False,
    entity_to_id=training.entity_to_id,     ## training
    relation_to_id=training.relation_to_id,)

validation = TriplesFactory.from_path(
    path=VALIDATE_PATH,
    create_inverse_triples=False,
    entity_to_id=training.entity_to_id,     ## training
    relation_to_id=training.relation_to_id,)


from pykeen.pipeline import pipeline
import torch
import pystow

directory = pystow.join('pykeen', 'tmp')
pipeline_result = pipeline(
    training=training,
    testing=testing,
    validation=validation,
    model='TransE',
    training_kwargs=dict(num_epochs=1, use_tqdm=False),
)
pipeline_result.save_to_directory(directory)

### just so we're not tempted to reuse this in the rest of the example
del pipeline_result

model = torch.load(directory / 'trained_model.pkl')

### now, just access the training data from the dataset and do the rest the same
list_idx = training.entities_to_ids(['egypt'])
idx = torch.as_tensor(list_idx, device=model.device)
entity_embeddings = model.entity_embeddings
egypt_embedding = entity_embeddings(idx).detach()[0]
egypt_embedding
```
Above is the code that I used my own local dataset to solve this problem. 
If you have better suggestions to solve this problem, I hope to get your help.
Thanks!

The pipeline provides the flexibility to use either a Dataset class or bring your own triples factories. If you're just using the vanilla Nations dataset, you can always specify you want inverse triples if you pre-instantiate it (this could cut out some code you wrote since it's very standard). Otherwise, I think what you wrote looks fine (besides the small number of epochs!)

If you're happy with the solution to your original issue, please close this ticket. You can also feel free to open another later if there's anything else :)

> The pipeline provides the flexibility to use either a Dataset class or bring your own triples factories. If you're just using the vanilla Nations dataset, you can always specify you want inverse triples if you pre-instantiate it (this could cut out some code you wrote since it's very standard). Otherwise, I think what you wrote looks fine (besides the small number of epochs!)
> 
> If you're happy with the solution to your original issue, please close this ticket. You can also feel free to open another later if there's anything else :)

Thanks！
This vanilla Nations dataset  is the dataset that I'm using as an example.
I'm actually using a larger dataset.
Do I use the following code the same way I defined the TriplesFactory?

```
from pykeen.datasets.base import PathDataset

# DBpedia
TRAIN_PATH = "./DBpedia/train.tsv"
TEST_PATH = "./DBpedia/test.tsv"
VALIDATE_PATH = "./DBpedia/valid.tsv"

# nations
# TRAIN_PATH = "./nations/train.txt"
# TEST_PATH = "./nations/test.txt"
# VALIDATE_PATH = "./nations/valid.txt"

class DBpedia(PathDataset):
    def __init__(self, create_inverse_triples: bool = False, **kwargs):
       # Initialize the DBpedia dataset.
        super().__init__(
            training_path=TRAIN_PATH,
            testing_path=TEST_PATH,
            validation_path=VALIDATE_PATH,
            create_inverse_triples=False,
            **kwargs,
        )

pipeline_result = pipeline(
    dataset="DBpedia",
    #training=training,
    #testing=testing,
    #validation=validation,
    model='TransE',
    training_kwargs=dict(num_epochs=1, use_tqdm=False),
)
```