1

How do i get an embedding for the whole sentence from huggingface's feature extraction pipeline?

I understand how to get the features for each token (below) but how do i get the overall features for the sentence as a whole?

feature_extraction = pipeline('feature-extraction', model="distilroberta-base", tokenizer="distilroberta-base")
features = feature_extraction("i am sentence")
stackoverflowuser2010
  • 29,060
  • 31
  • 142
  • 184
user3472360
  • 321
  • 3
  • 17

2 Answers2

1

To explain more on the comment that I have put under stackoverflowuser2010's answer, I will use "barebone" models, but the behavior is the same with the pipeline component.

BERT and derived models (including DistilRoberta, which is the model you are using in the pipeline) agenerally indicate the start and end of a sentence with special tokens (mostly denoted as [CLS] for the first token) that usually are the easiest way of making predictions/generating embeddings over the entire sequence. There is a discussion within the community about which method is superior (see also a more detailed answer by stackoverflowuser2010 here), however, if you simply want a "quick" solution, then taking the [CLS] token is certainly a valid strategy.

Now, while the documentation of the FeatureExtractionPipeline isn't very clear, in your example we can easily compare the outputs, specifically their lengths, with a direct model call:

from transformers import pipeline, AutoTokenizer

# direct encoding of the sample sentence
tokenizer = AutoTokenizer.from_pretrained('distilroberta-base')
encoded_seq = tokenizer.encode("i am sentence")

# your approach
feature_extraction = pipeline('feature-extraction', model="distilroberta-base", tokenizer="distilroberta-base")
features = feature_extraction("i am sentence")

# Compare lengths of outputs
print(len(encoded_seq)) # 5
# Note that the output has a weird list output that requires to index with 0.
print(len(features[0])) # 5

When inspecting the content of encoded_seq, you will notice that the first token is indexed with 0, denoting the beginning-of-sequence token (in our case, the embedding token). Since the output lengths are the same, you could then simply access a preliminary sentence embedding by doing something like

sentence_embedding = features[0][0]

dennlinger
  • 6,287
  • 1
  • 23
  • 43
  • 2
    Note that the embedding for the `[CLS]` token will be gibberish unless you fine-tune the model on a downstream task. I hypothesize that if you pool over the token embeddings as I suggested in my answer, then the resulting sentence embedding will have meaning without additional fine-tuning. The reason I bring this notion up is that the out-of-the-box `pipeline` classes do not provide an API for fine-tuning the underlying model. – stackoverflowuser2010 Nov 08 '20 at 21:13
  • Makes sense, also a good point! Would you suggest to only pool over actual tokens (since `[CLS]` and `[EOS]` are also tokens in the sample input on their own), or include the entire sequence regardless? – dennlinger Nov 09 '20 at 09:11
  • 1
    I would train on a downstream task to get good sentence embeddings. Using the NLI task seems to be the current best practice for doing so. I've been getting good empirical results by pooling over all the tokens, including subtokens (`##foo`) and special tokens (`[CLS]`, `[SEP]`), although I'd like to explore alternatives in the future. – stackoverflowuser2010 Nov 09 '20 at 19:23
  • I can't edit it cause the edition has less than 6 characters, but the `AutoTokenizer` has an invalid base name, it should be 'distilroberta-base'. –  Dec 30 '20 at 01:45
  • Good catch, I updated the answer! – dennlinger Dec 30 '20 at 08:49
0

If you have the embeddings for each token, you can create an overall sentence embedding by pooling (summarizing) over them. Note that if you have D-dimensional token embeddings, you should get a D-dimensional sentence embeddings through one of these approaches:

  1. Compute the mean over all token embeddings.

  2. Compute the max of each of the D-dimensions over all the token embeddings.

stackoverflowuser2010
  • 29,060
  • 31
  • 142
  • 184
  • I don't think this is necessarily the (only) correct approach. Ideally, you can simply use the embedding of the `[CLS]` token, which should act as an embedding layer. I'll try to post an answer of how to acess this via the `pipeline` feature. While I agree that averaging is also a valid approach, especially when pre-training or fine-tuning, most approaches out there utilize the `[CLS]` token and not the general token embeddings. – dennlinger Nov 06 '20 at 11:13