Add text embedding serving

[coderrg]

Resolves #206

[jonatanklosko]

@coderrg awesome thanks, just one nitpick from me!

@seanmor5 you've been using embeddings, so please check if this matches what you'd expect :D

[seanmor5]

The only thing I would add is an option to include a function which transforms the output attribute as a part of the serving. I'm not sure HF has this, but it seems common with sentence transformer embeddings

[rajrajhans]

Thanks for the great work!

A beginner question about this - if someone could help.

This line - encoder.(params, input)[output_attribute] mandatorily extracts output_attribute from encoder.(params, input).

However, if I'm trying to get CLIP text embedding with textual_projection, then my model would look something like -

{:ok, %{model: text_model, params: text_params, spec: text_spec}} =
        Bumblebee.load_model({:hf, clip_model_name},
          module: Bumblebee.Text.ClipText,
          architecture: :base
        )

my_model = text_model
           |> Axon.nx(& &1.pooled_state)
           |> Axon.dense(512, use_bias: false, name: "text_projection")

Here, the output of my_model would be the actual Nx tensor output that I would like. But, since the text embedding always tries to extract output_attribute, I will have to add another layer to shift the actual output inside an object

my_model_2 =
        text_model
        |> Axon.nx(& &1.pooled_state)
        |> Axon.dense(512, use_bias: false, name: "text_projection")
        |> Axon.nx(fn x -> %{my_output_attribute: x} end)

And then pass :my_output_attribute as the value for output_attribute option.

1. Is there another way around this? Please correct me if I'm wrong.
2. Maybe having output_attribute optional here would be a good way out?

[coderrg]

I appreciate the suggestions @jonatanklosko and @seanmor5 — I've implemented those changes! @seanmor5 I know folks will sometimes apply L2 normalization to their embeddings, so I added support for it; let me know if there were other functions you had in mind. And @rajrajhans, thank you for bringing that up. For models like CLIP that have a projection head, it would be helpful to have the option to directly retrieve the model output.

My guess is that for most models, the pooled state (as an attribute of the output) would be used as the embedding, so perhaps it would be best to add a non-default option :none for output_attribute to account for the projection head case (e.g., CLIP). I've gone ahead and done this for now, but I can also make output_attribute optional altogether if that would be preferred. Curious to hear what other folks think.

[trodrigu]

This is looking great @coderrg! I was curious about the postprocessing and mean pooling. Should this be a postprocessing function or is this out of scope?

[coderrg]

Thanks @trodrigu — that makes sense. I've added support for mean_pooling as an embedding function based on your example since it seems useful for sentence transformers. Does my implementation line up with what you'd expect?

Also, since mean pooling and other functions applied to output embeddings are not mutually exclusive, I changed embedding_function to embedding_functions so that it takes a list of functions to be applied to the output embedding in order. The supported functions are now :l2_normalization and :mean_pooling.

[jonatanklosko]

Awesome, thanks for the feedback everyone!

@coderrg I left a couple more comments, but it's looking great :D

[coderrg]

Thanks @jonatanklosko, I've implemented your suggestions. I also added nil as the default option for :output_pool and :embedding_processor — please let me know if I misunderstood that suggestion, though.

[jonatanklosko]

Perfect, a couple final comment and we should be good to go!

[jonatanklosko]

Awesome, thanks!