[fix] Use return_dict=True in Transformer; improve how all_layer_embeddings are determined

[tomaarsen]

Fixes #3318

Hello!

Pull Request overview

• Improve how all_layer_embeddings are determined

Details

In short, instead of return_dict=False where we just get a tuple without any keys, we now use return_dict=True, allowing us to see the output names. transformers seems to have a pretty strong convention to use hidden_states for the output_hidden_states, so this should be fairly safe. Beyond that, I'm still using [0] indexing for the token embeddings (also known as last_hidden_state).

My intention is not to introduce any backwards breaking here, but there's always a risk. Let me know if your code breaks because of this!

• Tom Aarsen

[Copilot]

Copilot reviewed 1 out of 1 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (1)

sentence_transformers/models/Transformer.py:443

• [nitpick] Using indexing with [0] to extract token embeddings may be unclear now that a dictionary is returned; consider using outputs['last_hidden_state'] to improve code clarity and maintain consistency with typical transformers output.

token_embeddings = outputs[0]