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)

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