Skip to content

Add support for CLS token pooling in text embedding #385

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Aug 6, 2024
Merged

Add support for CLS token pooling in text embedding #385

Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
78dc2e3
add suport for cls token pooling
nyo16 Aug 2, 2024
32dfdc3
removing test
nyo16 Aug 2, 2024
3535141
simplifying Nx slice
nyo16 Aug 2, 2024
574128c
fixing formating issues.
nyo16 Aug 2, 2024
61fe5d2
documentation added to main module. Simplifying logic for check the a…
nyo16 Aug 6, 2024
5ce9de4
Update lib/bumblebee/text.ex
jonatanklosko Aug 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion lib/bumblebee/text/text_embedding.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,18 @@ defmodule Bumblebee.Text.TextEmbedding do
nil ->
output

:cls_token_pooling ->
case Nx.rank(output) do
3 ->
# Assuming CLS token is always at the first position
Nx.take(output, 0, axis: 1)

rank ->
raise ArgumentError,
"expected the output tensor to have rank 3 to apply :cls pooling, got: #{rank}." <>
" You should either disable pooling or pick a different output using :output_attribute"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe for any pooling we expect 3, because we reduce n tokens into 1. Now that we have more types, we can move the check before the case, like this:

if output_pool != nil and Nx.rank(output) != 3 do
  raise ...
end

We can use the message from the other clause!

end

:mean_pooling ->
case Nx.rank(output) do
3 ->
Expand All @@ -81,7 +93,7 @@ defmodule Bumblebee.Text.TextEmbedding do

other ->
raise ArgumentError,
"expected :output_pool to be one of nil or :mean_pooling, got: #{inspect(other)}"
"expected :output_pool to be one of :cls_token_pooling, :mean_pooling or nil, got: #{inspect(other)}"
end

output =
Expand Down
Loading