Add zero-shot vision

lib/bumblebee/vision.ex

@@ -214,4 +214,14 @@ defmodule Bumblebee.Vision do
         ) :: Nx.Serving.t()
   defdelegate image_embedding(model_info, featurizer, opts \\ []),
     to: Bumblebee.Vision.ImageEmbedding
+
+  @spec zero_shot_classification(
+          Bumblebee.model_info(),
+          Bumblebee.Featurizer.t(),
+          Bumblebee.Tokenizer.t(),
+          list(),
+          keyword()
+        ) :: Nx.Serving.t()
+  defdelegate zero_shot_classification(model_info, featurizer, tokenizer, labels, opts \\ []),
+    to: Bumblebee.Vision.ZeroShotClassification
 end

lib/bumblebee/vision/zero_shot_classification.ex

@@ -0,0 +1,89 @@
+defmodule Bumblebee.Vision.ZeroShotClassification do
+  @moduledoc false
+
+  alias Bumblebee.Utils
+  alias Bumblebee.Shared
+
+  def zero_shot_classification(model_info, featurizer, tokenizer, labels, opts \\ []) do
+    %{model: model, params: params, spec: spec} = model_info
+    Shared.validate_architecture!(spec, :base)
+
+    opts =
+      Keyword.validate!(opts, [
+        :compile,
+        hypothesis_template: &default_hypothesis_template/1,
+        defn_options: [],
+        preallocate_params: false
+      ])
+
+    hypothesis_template = opts[:hypothesis_template]
+    preallocate_params = opts[:preallocate_params]
+    defn_options = opts[:defn_options]
+
+    hypotheses = Enum.map(labels, hypothesis_template)
+
+    tokenized_hypotheses =
+      Nx.with_default_backend(Nx.BinaryBackend, fn ->
+        Bumblebee.apply_tokenizer(tokenizer, hypotheses, return_token_type_ids: false)
+      end)
+
+    compile =
+      if compile = opts[:compile] do
+        compile
+        |> Keyword.validate!([:batch_size])
+        |> Shared.require_options!([:batch_size])
+      end
+
+    batch_size = compile[:batch_size]
+
+    {_init_fun, predict_fun} = Axon.build(model)
+
+    logits_fun = fn params, input ->
+      input =
+        featurizer
+        |> Bumblebee.Featurizer.process_batch(input)
+        |> Map.merge(tokenized_hypotheses)
+
+      %{logits_per_image: logits_per_image} = predict_fun.(params, input)
+
+      logits_per_image
+    end
+
+    Nx.Serving.new(
+      fn defn_options ->
+        params = Shared.maybe_preallocate(params, preallocate_params, defn_options)
+
+        logits_fun =
+          Shared.compile_or_jit(logits_fun, defn_options, compile != nil, fn ->
+            inputs = Bumblebee.Featurizer.batch_template(featurizer, batch_size)
+            [params, inputs]
+          end)
+
+        fn inputs ->
+          inputs = Shared.maybe_pad(inputs, batch_size)
+          logits = logits_fun.(params, inputs)
+          Axon.Activations.softmax(logits)
+        end
+      end,
+      defn_options
+    )
+    |> Nx.Serving.batch_size(batch_size)
+    |> Nx.Serving.client_preprocessing(fn input ->
+      {images, multi?} = Shared.validate_serving_input!(input, &Shared.validate_image/1)
+      inputs = Bumblebee.Featurizer.process_input(featurizer, images)
+      {Nx.Batch.concatenate([inputs]), multi?}
+    end)
+    |> Nx.Serving.client_postprocessing(fn {scores, _metadata}, multi? ->
+      scores
+      |> Utils.Nx.to_list()
+      |> Enum.map(fn scores_for_batch ->
+        Enum.zip_with(scores_for_batch, labels, fn score, label ->
+          %{score: score, label: label}
+        end)
+      end)
+      |> Shared.normalize_output(multi?)
+    end)
+  end
+
+  defp default_hypothesis_template(label), do: "This is a photo of #{label}."
+end