Add state to logits processing (#425)

Co-authored-by: Joel Koch <joel@bitcrowd.net>
Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

Author: 3 people

lib/bumblebee.ex

@@ -1089,6 +1089,36 @@ defmodule Bumblebee do
     end
   end
 
+  @doc """
+  Initializes state for a new logits processor.
+
+  Returns `state`, which is an opaque `Nx.Container`, and it is then
+  passed to and returned from `process/4`.
+  """
+  @doc type: :logits_processor
+  @spec logits_processor_init(
+          Bumblebee.LogitsProcessor.t(),
+          context :: Bumblebee.LogitsProcessor.init_context()
+        ) :: Bumblebee.LogitsProcessor.state()
+  def logits_processor_init(%module{} = logits_processor, context) do
+    module.init(logits_processor, context)
+  end
+
+  @doc """
+  Processes logits, applying specific rules. Receives context, state and
+  logits, and returns updated logits and state.
+  """
+  @doc type: :logits_processor
+  @spec logits_processor_process(
+          Bumblebee.LogitsProcessor.t(),
+          Bumblebee.LogitsProcessor.state(),
+          logits :: Nx.Tensor.t(),
+          context :: Bumblebee.LogitsProcessor.process_context()
+        ) :: {Bumblebee.LogitsProcessor.state(), logits :: Nx.Tensor.t()}
+  def logits_processor_process(%module{} = logits_processor, state, logits, context) do
+    module.process(logits_processor, state, logits, context)
+  end
+
   @doc """
   Initializes state for a new scheduler loop.
 

lib/bumblebee/logits_processor.ex

@@ -0,0 +1,46 @@
+defmodule Bumblebee.LogitsProcessor do
+  @moduledoc """
+  An interface for configuring and using logits processors.
+
+  Logits processors are used during autoregressive generation to modify
+  predicted scores at each generation step. This allows for applying
+  certain rules to the model output to control which tokens are picked
+  at each generation step, and which are not.
+
+  Every module implementing this behaviour is expected to also define
+  a configuration struct.
+  """
+
+  @type t :: Bumblebee.Configurable.t()
+
+  @type state :: Nx.Container.t()
+
+  @type process_context :: %{
+          sequence: Nx.Tensor.t(),
+          length: Nx.Tensor.t(),
+          input_length: Nx.Tensor.t()
+        }
+
+  @type init_context :: %{}
+
+  @doc """
+  Initializes state for a new logits processor.
+
+  Returns `state`, which is an opaque `Nx.Container`, and it is then
+  passed to and returned from `process/2`.
+
+  Oftentimes logits processors are stateless, in which case this
+  function can return an empty container, such as `{}`.
+  """
+  @callback init(t(), init_context()) :: state()
+
+  @doc """
+  Processes logits, applying specific rules.
+  """
+  @callback process(
+              t(),
+              state(),
+              logits :: Nx.Tensor.t(),
+              context :: process_context()
+            ) :: {state :: map(), logits :: Nx.Tensor.t()}
+end