diff --git a/lib/bumblebee/text/gemma3_text.ex b/lib/bumblebee/text/gemma3_text.ex index 19301315..da0e72e3 100644 --- a/lib/bumblebee/text/gemma3_text.ex +++ b/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 diff --git a/test/bumblebee/text/gemma3_text_test.exs b/test/bumblebee/text/gemma3_text_test.exs index f1b4de39..8ac7e67e 100644 --- a/test/bumblebee/text/gemma3_text_test.exs +++ b/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(