Fix sliding window size for bidirectional Gemma 3

lib/bumblebee/text/gemma3_text.ex

@@ -80,7 +80,7 @@ defmodule Bumblebee.Text.Gemma3Text do
       attention_window_size: [
         default: 4096,
         doc:
-          "the number of tokens each token can attend to in the sliding attention window, including itself. Only used for `:sliding_attention` layers"
+          "the number of tokens each token can attend to in the sliding attention window, including itself. With bidirectional attention, the window is centered on the token, spanning half the size on each side. Only used for `:sliding_attention` layers"
       ],
       layer_types: [
         default: nil,
@@ -395,8 +395,14 @@ defmodule Bumblebee.Text.Gemma3Text do
               {nil, spec.rotary_embedding_base, spec.rotary_embedding_scaling_strategy}
 
             :sliding_attention ->
-              {{spec.attention_window_size - 1, spec.attention_window_size - 1},
-               spec.rotary_embedding_base_local, nil}
+              # With bidirectional attention, the window spans both sides
+              # see https://github.com/huggingface/transformers/blob/v5.20.2/src/transformers/models/gemma3/configuration_gemma3.py#L102-L103
+              size =
+                if spec.use_bidirectional_attention,
+                  do: div(spec.attention_window_size, 2),
+                  else: spec.attention_window_size - 1
+
+              {{size, size}, spec.rotary_embedding_base_local, nil}
           end
 
         shortcut = hidden_state

test/bumblebee/text/gemma3_text_test.exs

@@ -58,6 +58,35 @@ defmodule Bumblebee.Text.Gemma3TextTest do
     )
   end
 
+  test ":base with bidirectional attention and sliding window" do
+    assert {:ok, %{model: model, params: params, spec: spec}} =
+             Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-Gemma3TextModel"},
+               spec_overrides: [use_bidirectional_attention: true, attention_window_size: 4]
+             )
+
+    assert %Bumblebee.Text.Gemma3Text{
+             architecture: :base,
+             use_bidirectional_attention: true
+           } = spec
+
+    inputs = %{
+      "input_ids" => Nx.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]),
+      "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
+    }
+
+    outputs = Axon.predict(model, params, inputs)
+
+    assert Nx.shape(outputs.hidden_state) == {1, 10, 32}
+
+    assert_all_close(
+      outputs.hidden_state[[.., 1..3, 1..3]],
+      Nx.tensor([
+        [[-0.6205, 1.8721, 1.7530], [-0.3245, 0.9898, 3.1781], [0.0932, 0.4541, 1.8766]]
+      ]),
+      atol: 1.0e-4
+    )
+  end
+
   test ":for_sequence_classification" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(