Embeddings V2

[KulikDM]

Hi,

This commit addresses #111
I have made two PRs where the functionality is slightly different, but they essentially both do the exact same thing.
Please pick which one you think is best for the package.
V2 adds the get_embeddings to each estimator.

Benefit is easier to find API, downside is more duplicate code.

Also added the appropriate tests.

Hope this helps!

[iivalchev]

Perhaps I am missing something but I don't see how you are getting the train/test embeddings for X. Shouldn't the iter_outputs method be passed a new param that would invoke the model with only_return_standard_out=False?

[noahho]

Thanks so much for looking into this! Unfortunately the change does not actually retrieve the embeddings now, see the review comments.

[KulikDM]

Oh shoot!

Sorry, I saw in the transformer.py the code block

output_decoded["train_embeddings"] = train_encoder_out
output_decoded["test_embeddings"] = test_encoder_out

And traced it back to the regressor and classifier.

Would adding the arg only_return_standard_out=False do the trick?
If so, I can make the adjustment to this version as well as updating the tests.

[iivalchev]

Oh shoot!

Sorry, I saw in the transformer.py the code block

output_decoded["train_embeddings"] = train_encoder_out
output_decoded["test_embeddings"] = test_encoder_out

And traced it back to the regressor and classifier.

Would adding the arg only_return_standard_out=False do the trick? If so, I can make the adjustment to this version as well as updating the tests.

Looks that way, but you need to indicate this through the tabpfn.inference.InferenceEngine.iter_outputs with additional param to it, say include_embbedings_in_output wich defaults to False and only get_embeddings should set it to True. One more thing you should change how you handle what would then be returned from forward pass of the model as it would be a dict not just a single tensor.

[KulikDM]

Hi @iivalchev

Thanks for all the feedback!

The latest commit tries to incorporate all the suggestions mentioned above.

• Single function in utils that is called by the classifier and the regressor
• Pass only_return_standard_out as a new arg to the inference.iter_outputs (default set to True like before, embedding passes False)
• iter_ouput checks output type, and if type dict, it extracts the test_embeddings (not sure if this should be train was following Embeddings Functionality #111)
• Tests have been updated

Everything was working as expected on my side; hope this change is acceptable.

[noahho]

This change is a bit hacky for the application, ideally, we would retrieve the entire dict at this step and in the next step decide which entry to select. That would change the return type of this function.
Alternatively we could pass a parameter return_field = 'test_embeddings' which selects whcih field to select and which is linked to the only_return_standard_out. The first option would be cleaner though?

[noahho]

Great! I think this should be working already, I just had some comments to make it more maintainable.

[KulikDM]

Hi @noahho

Should I make the changes to rather pass the dictionary to the _get_embeddings function and add an additional arg for the user to select which information they would like to retrieve?

[noahho]

Hi @noahho

Should I make the changes to rather pass the dictionary to the _get_embeddings function and add an additional arg for the user to select which information they would like to retrieve?

I would let iter_outputs return the whole dictionary if not returning single out. Then Let get_embeddings_ select the test or train embeddings where the user specifies if would like to get test or train embedding. You think that makes sense?

[KulikDM]

Hi @noahho

Awesome, agreed!

I have added the additional arg data_source to either select the test (default) or train embeddings in the get_embeddings function.

[noahho]

Great! I just fixed the merge conflict and now CI should run.

[CLAassistant]

All committers have signed the CLA.

[KulikDM]

Hi @noahho

Seems like it is failing for python 3.9 but not 3.12
The issue is an infinite showing up.
Tested locally and it seems to fail due to passing a random_state value during regressor fit. Removing the arg resolved the issue. Will add this change later today :)

[hzhz2020]

Dear Authors

Thanks for this great work! The performance in the paper is amazing and I have heard many good things about the work. I am eager to try it myself. I'm wondering do you have an expected timeline when this feature (obtaining embedding) would be available?

Thanks!

[noahho]

PR looks great, thank you for the contribution @KulikDM !

[lincj1994]

@KulikDM Hi. Thanks for the great work. I'm still confused of how to get embeddings from a fitted model, such as the clf in the code example below. Could you please provide some additional guidance?
Thanks.

import pandas as pd
import numpy as np
from tabpfn import TabPFNClassifier
from tabpfn_extensions.post_hoc_ensembles.sklearn_interface import AutoTabPFNClassifier
from sklearn.model_selection import train_test_split


df = pd.read_csv('df2.csv')
X_features = list(df.drop(['resp_crpr', 'Tumor_Sample_Barcode'], axis=1))
y_features = ['resp_crpr']
clf = AutoTabPFNClassifier(device='auto')


args = Namespace(
    seed = 444444,
    test_size = 0.3
)
X_train, X_test, y_train, y_test = train_test_split(df[X_features].values.reshape(-1, 63), df[y_features], 
                                                test_size=args.test_size, random_state=args.seed)
clf.fit(X_train, y_train)