Extract the embedding from a specific layer of MarianModel

Question

I am using using MarianModel from the hub of HuggingFace for a translation task. Now I want to extract the embedding from the output of the last MarianEncoderLayer layer. Specifically, given a text, I want to get the middle embedding of the text for another task. Here is what I did:

text = 'hello how are you?'
inputs = tokenizer(text,max_length=512,truncation=True,return_tensors='pt',padding=True)
model.model.encoder.layers[5](**inputs)

But it throws the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\dohuut\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'input_ids'

What is the proper way to extract this embedding?

Answer 1

You are passing discrete tokens as input to an attention layer that expects vectors of real numbers as inputs. The error you are getting tells you that the layer does not expect an input parameter called input_ids, which is what you are giving to it.

Instead, you could use the whole model, passing output_hidden_states=True as a parameter (see the documentation). Then, you could extract the output layer you want from the obtained hidden states (encoder_hidden_states).

The structure of encoder_hidden_states is described in the documentation:

encoder_hidden_states (tuple(torch.FloatTensor), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of torch.FloatTensor (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape (batch_size, sequence_length, hidden_size).

Hidden-states of the encoder at the output of each layer plus the optional initial embedding outputs.