search

xenova / transformers.js

State-of-the-art Machine Learning for the web. Run 🤗 Transformers directly in your browser, with no need for a server!
https://huggingface.co/docs/transformers.js
Apache License 2.0
11.04k stars 678 forks source link

Update pipelines.js to allow for `token_embeddings` as well #770

Closed NikhilVerma closed 4 months ago

NikhilVerma commented 4 months ago

In recent examples of optimum pipeline export the feature extraction pipelines have their output state as token_embeddings instead of last_hidden_state.

So we should support this as well.

xenova commented 4 months ago

Thanks for this! I'm pretty sure this is only for sentence-transformers models exported via Optimum (see here), but a good update regardless!

xenova commented 4 months ago

Can you give an example of a model which has this as an output name?

NikhilVerma commented 4 months ago

@xenova Sorry I should have provided. this is the command I ran

optimum-cli export onnx --model sentence-transformers/all-MiniLM-L6-v2 ./sbert/

I am using the pnpm patch-package ability to do this and it works quite well!

xenova commented 4 months ago

Thanks! Since the exported model returns two outputs (token_embeddings and sentence_embedding), I think it will be a good idea to try get the sentence_embedding if it is present and the user does not specify a pooling method. Otherwise, if pooling is set, we compute it the normal way. We should also allow the user to normalize their sentence embeddings with the normalize option.

Is this something you'd like to add to the PR?

HuggingFaceDocBuilderDev commented 4 months ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

xenova commented 4 months ago

^^ Will merge for now, but could be a good follow-up PR. Let me know if you're interested!