introduced types for init_context and process_context

Author: xhr15

lib/bumblebee.ex

@@ -1092,7 +1092,7 @@ defmodule Bumblebee do
   @doc type: :logits_processor
   @spec logits_processor_init(
           Bumblebee.LogitsProcessor.t(),
-          context :: term()
+          context :: Bumblebee.LogitsProcessor.init_context()
         ) :: Bumblebee.LogitsProcessor.state()
   def logits_processor_init(%module{} = logits_processor, context) do
     module.init(logits_processor, context)
@@ -1107,7 +1107,7 @@ defmodule Bumblebee do
           Bumblebee.LogitsProcessor.t(),
           Bumblebee.LogitsProcessor.state(),
           logits :: Nx.Tensor.t(),
-          context :: term()
+          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)

lib/bumblebee/logits_processor.ex

@@ -15,6 +15,14 @@ defmodule Bumblebee.LogitsProcessor do
 
   @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.
 
@@ -24,7 +32,7 @@ defmodule Bumblebee.LogitsProcessor do
   Oftentimes logits processors are stateless, in which case this
   function can return an empty container, such as `{}`.
   """
-  @callback init(t(), any()) :: state()
+  @callback init(t(), init_context()) :: state()
 
   @doc """
   Processes logits, applying specific rules.
@@ -33,6 +41,6 @@ defmodule Bumblebee.LogitsProcessor do
               t(),
               state(),
               logits :: Nx.Tensor.t(),
-              context :: term()
+              context :: process_context()
             ) :: {state :: map(), logits :: Nx.Tensor.t()}
 end

lib/bumblebee/text/generation/stateless_logits_processor.ex

@@ -19,12 +19,12 @@ defmodule Bumblebee.Text.Generation.StatelessLogitsProcessor do
   end
 
   @impl Bumblebee.LogitsProcessor
-  def init(_logits_processor, _context) do
+  def init(_logits_processor, _init_context) do
     %{}
   end
 
   @impl Bumblebee.LogitsProcessor
-  def process(logits_processor, state, logits, context) do
-    {state, logits_processor.fun.(logits, context)}
+  def process(logits_processor, state, logits, process_context) do
+    {state, logits_processor.fun.(logits, process_context)}
   end
 end

test/bumblebee/text/generation_test.exs

@@ -220,19 +220,19 @@ defmodule Bumblebee.Text.GenerationTest do
     end
 
     @impl Bumblebee.LogitsProcessor
-    def init(logits_processor, context) do
+    def init(logits_processor, init_context) do
       initial_enforced_token_id = Nx.tensor([logits_processor.initial_enforced_token_id])
 
       [initial_enforced_batch_token_id, _sequence] =
-        Nx.broadcast_vectors([initial_enforced_token_id, context.sequence])
+        Nx.broadcast_vectors([initial_enforced_token_id, init_context.sequence])
 
       %{
         next_enforced_token_id: initial_enforced_batch_token_id
       }
     end
 
     @impl Bumblebee.LogitsProcessor
-    def process(_logits_processor, state, logits, _context) do
+    def process(_logits_processor, state, logits, _process_context) do
       next_enforced_token_id = state.next_enforced_token_id
 
       logits = enforce_token(logits, next_enforced_token_id)