Support for T5Gemma2 Encoder

[contrebande-labs]

Hi guys,

For the fun of it, I tried to load a SentenceTransformers on the google/t5gemma-2-270m-270m configs and weights with this line: st_pytorch_model = SentenceTransformer(input_model_root_path, backend="torch", device="xpu", trust_remote_code=True, local_files_only=True, model_kwargs={"dtype": torch.float32}).eval(). I'm getting this trace:

Loading pytorch model: /media/user/500gb1/models/huggingface/google/t5gemma-2-270m-270m
No sentence-transformers model found with name /media/user/500gb1/models/huggingface/google/t5gemma-2-270m-270m. Creating a new one with mean pooling.
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 911/911 [00:13<00:00, 66.90it/s, Materializing param=encoder.vision_tower.vision_model.post_layernorm.weight]
Traceback (most recent call last):
  File "/home/user/Documents/ir-sts/./t5gemma2.py", line 17, in <module>
    st_pytorch_model = SentenceTransformer(input_model_root_path, backend="torch", device="xpu", trust_remote_code=True, local_files_only=True, model_kwargs={"dtype": torch.float32}).eval()
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/venv/optimum-onnx-git/lib/python3.12/site-packages/sentence_transformers/SentenceTransformer.py", line 339, in __init__
    modules = self._load_auto_model(
              ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/venv/optimum-onnx-git/lib/python3.12/site-packages/sentence_transformers/SentenceTransformer.py", line 2120, in _load_auto_model
    pooling_model = Pooling(transformer_model.get_word_embedding_dimension(), "mean")
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/venv/optimum-onnx-git/lib/python3.12/site-packages/sentence_transformers/models/Transformer.py", line 288, in get_word_embedding_dimension
    return self.auto_model.config.hidden_size
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/venv/optimum-onnx-git/lib/python3.12/site-packages/transformers/configuration_utils.py", line 200, in __getattribute__
    return super().__getattribute__(key)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'T5Gemma2Config' object has no attribute 'hidden_size'

I don't know if it's a good idea in general (I don't even know what they mean by tied embedding, but is there a quick hack to fix the 'T5Gemma2Config' object has no attribute 'hidden_size' error and be able to try the model as a sentence embedder? Thanks!

[tomaarsen]

Hello!

The prior T5 model had to be loaded via the T5EncoderModel (

sentence-transformers/sentence_transformers/models/Transformer.py

Line 217 in c19c540

def _load_t5_model(self, model_name_or_path: str, config: PretrainedConfig, cache_dir: str, **model_args) -> None:

), perhaps for the same reason that you're encountering now. Beyond that, Sentence Transformers is primarily interested in the (last) hidden state from the encoder, so using the decoder presumably doesn't help.

If I were you, I would subclass the sentence_transformers.models.Transformer and reimplement _load_model to load a T5 Gemma2 encoder if possible. Then, you can initialize the SentenceTransformer without the SentenceTransformer(model_name_or_path) shorthand, but instead using the SentenceTransformer(modules=[transformer, pooling, normalize]) format, see e.g. https://sbert.net/docs/sentence_transformer/usage/custom_models.html

After the holidays, I can try and push an update where the T5Gemma encoder is used automatically when a T5Gemma2 model is loaded, akin to what's already being done for T5 and MT5.

• Tom Aarsen

[contrebande-labs]

Hi Tom!

Thanks for taking the time to reply. I was trying to isolate the encoder from the decoder, but I didn't find an encode method anywhere in the HF documentation. The next step for me before I talked to you would have been to look into the HF-implemented T5Gemma2ForSequenceClassification and T5Gemma2ForTokenClassification to see if and how they isolate the encoder's last hidden states. But I'm also now going to peek into the code you mentioned in ST to load the encoder module separately and initialize a pooling and normalize module from scratch. I will report here.

Basically I'm trying to see how Deepmind came up with embeddinggemma as it's my favorite text embedder at the moment.

hanks again and happy Holidays!

-- Vincent

[contrebande-labs]

Hi! @tomaarsen !

So I had a look at the t5gemma2 modelling code in the transformers repo and the best bootstrap class I could find for what we want is T5Gemma2Encoder initialized with the encoder part of the config. Unfortunately, there does not seem to be a clean cut separation between the text encoder and the siglip head in there. So I think I should manually implement a new T5Gemma2TextEncoder subclass of sentence_transformers.models.Transformer without all the multimodal stuff. Also, can T5Gemma2TextEncoder be hosted in sentence-transformers somewhere instead of in transformers. I think I won't survive dipping my toes in this mess. If so, where would you like me to put it? I can put it all in a _load_t5gemma2_model method, but I think it might bloat the whole thing a little too much (but maybe not). Let me know what you think so I can fork, code, commit, push and PR. Hope, you're having a great holiday time!