How to use T5 for sentence embedding?

is there any way to use encoder part of T5 model for representation learning?

Hi @banucool
You can initialize the T5Model class and only forward pass through it’s encoder. The first element of the returned tuple is the final hidden states.

model = T5Model.from_pretrained("t5-small")
tok = T5Tokenizer.from_pretrained("t5-small")

enc = tok("some text", return_tensors="pt")

# forward pass through encoder only
output = model.encoder(
    input_ids=enc["input_ids"], 
    attention_mask=enc["attention_mask"], 
    return_dict=True
)
# get the final hidden states
emb = output.last_hidden_state

The shape of emb will be (batch_size, seq_len, hidden_size)

2 Likes

thanks a lot @valhalla :blush:

can we use pruned version of bert for feature extraction?does it make sense?

To clarify, the above code just returns the final hidden state of each token and not whole sentence embedding.
for sentence embedding you can try sentence-bert.
https://huggingface.co/sentence-transformers

Hi, I’m interested in using T5 to generate word embeddings. I tried the code supplied above. Unfortunately, got this error message:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-40-5f6e22d1ad1e> in <module>()
      1 model = T5Model.from_pretrained("t5-small")
----> 2 tok = T5Tokenizer.from_pretrained("t5-small")
      3 
      4 enc = tok("some text", return_tensors="pt")
      5 

TypeError: 'NoneType' object is not callable

Do you have any thoughts on resolving this error message?

Thank you in advance for your help. :slightly_smiling_face: