sentence-transformers

/

static-retrieval-mrl-en-v1

like 45

Follow

Sentence Transformers 2.41k

This model is not transformers compatible, I'm afraid. You can only run it with sentence transformers.

• Tom Aarsen

I saw that it was possible with transformers.js: https://github.com/huggingface/transformers.js/issues/1160
but i don't know how to put the model type in python transformers.
I wanna try late chunking since the sequence length is infinite, but I couldn't really figure out how to do it with sentence-transformers

model = SentenceTransformer(
    "tomaarsen/static-retrieval-mrl-en-v1",
    backend="onnx",
    device="cpu"
)

tokens = model.tokenize(sentences)
print(tokens)

print(model.encode("".join(sentences), output_value="token_embeddings", convert_to_tensor=True))

when i try something like this get

for token_emb, attention in zip(out_features[output_value], out_features["attention_mask"]):
                                    ~~~~~~~~~~~~^^^^^^^^^^^^^^
KeyError: 'token_embeddings'

so i guess token embeddings are not supported by this model?

Ah, understandable. Indeed, the StaticEmbedding module that is used in this model (as seen in modules.json) uses an EmbeddingBag, which grabs token embeddings and then applies mean pooling. This EmbeddingBag is equivalent to Embedding plus torch.mean, but a bit faster to run.

If you want to try late interaction with it, then I would recommend copying the StaticEmbedding into a new module locally, changing the EmbeddingBag to Embedding, updating the forward so that the features["token_embeddings"] = self.embedding(...), cloning this model, and updating the modules.json to e.g.:

[
  {
    "idx": 0,
    "name": "0",
    "path": "",
    "type": "my_custom_static_embedding_script.TokenStaticEmbedding"
  },
  {
    "idx": 1,
    "name": "1",
    "path": "1_Pooling",
    "type": "sentence_transformers.models.Pooling"
  }
]

• Tom Aarsen

Hey thanks for answering. I am kinda new to this so I am having some problems.
I copied the StaticEmbedding to a new module locally, renamed the class and changed the EmbeddingBag to embedding and updated the forward function to this:

  def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
        features["token_embeddings"] = self.embedding(features["input_ids"])
        return features

But I am having problems with the last part. I put the model.safetensors in my dir with tokenizer and populated modules.json with what you sent and created a directory: "1_Pooling" with this config

{
  "word_embedding_dimension": 1024,
  "pooling_mode_cls_token": false,
  "pooling_mode_mean_tokens": true,
  "pooling_mode_max_tokens": false,
  "pooling_mode_mean_sqrt_len_tokens": false
}

And when i run:

from sentence_transformers import SentenceTransformer
model = SentenceTransformer(
    "./",
    backend="onnx",
    device="cpu"
)

sentences = [
    "hello how are you",
    "whats up with you"
]

tokens = model.tokenize(sentences)
print(tokens)

print(model.encode("".join(sentences), output_value="token_embeddings", convert_to_tensor=True))

the tokenize works but the encoding throws this:

Traceback (most recent call last):
  File "/redacted/stuff/main.py", line 16, in <module>
    print(model.encode("".join(sentences), output_value="token_embeddings", convert_to_tensor=True))
          ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/redacted/.venv/lib/python3.13/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
  File "/redacted/.venv/lib/python3.13/site-packages/sentence_transformers/SentenceTransformer.py", line 1051, in encode
    out_features = self.forward(features, **kwargs)
  File "/redacted/.venv/lib/python3.13/site-packages/sentence_transformers/SentenceTransformer.py", line 1132, in forward
    input = module(input, **module_kwargs)
  File "/redacted/.venv/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/redacted/.venv/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/redacted/.venv/lib/python3.13/site-packages/sentence_transformers/models/Pooling.py", line 239, in forward
    output_vector = torch.cat(output_vectors, 1)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

Sorry I quite dont understand what I am doing