Custom pooling methods

[ahmedkooli]

Hey ! I have a question about pooling methods. While running a model and changing the --pooling-method (cls, mean and auto), it seems that the embeddings results are not changing. Is that normal?
Also, is there a way to choose a custom pooling method from python code?

The model I'm using is https://huggingface.co/dunzhang/stella_en_400M_v5

Thanks a lot, don't hesitate if you need more details.

[michaelfeil]

@ahmedkooli Custom pooling needs to be enabled for each embedding backend (torch, ctranslate, optimum). With your model, it is defined to the model on how it was trained. You can find a the information for your model here: https://huggingface.co/dunzhang/stella_en_400M_v5/blob/main/modules.json and here https://huggingface.co/dunzhang/stella_en_400M_v5/blob/main/1_Pooling/config.json Overwriting the default is not possible atm.

[ahmedkooli]

Thanks for the quick answer @michaelfeil ! I'm still not getting why results are different from sentence_transformers embeddings. Here's an simple example:

import requests
from sentence_transformers import SentenceTransformer
import torch

model_name = "dunzhang/stella_en_400M_v5"
sentences = ["Hey I'm using infinity."]
 
# SentenceTransformers embeddings
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SentenceTransformer(model_name, device=dev, trust_remote_code=True)
embeddings_st = model.encode(sentences)

# Infinity embeddings
answer = requests.post("http://localhost:7997/embeddings",
                       json={"model": model_name, "input": sentences})
embeddings_inf = torch.tensor([v["embedding"] for v in answer.json()["data"]])

Which results in 2 very different embeddings. My deployment looks like this:

docker run -it --gpus all \
    -v $(PWD)/data:/app/.cache \
    -p 7997:7997 \
    --name infinity \
    infinity-image \
    v2 \
    --model-id dunzhang/stella_en_400M_v5 \
    --pooling-method auto \
    --port 7997 \
    --trust-remote-code

However, when using sentence-transformers/all-MiniLM-L6-v2, embeddings are exactly the same.

Thanks again for the help!