How to use T5 for sentence embedding?

Research
2

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)

1 Like