From b5200a33f1237b2a60ee92cfcf1df5221706c763 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Fri, 5 Apr 2024 03:51:53 +1100
Subject: [PATCH] Gemma inference fix

---
 pyproject.toml          | 2 +-
 unsloth/models/gemma.py | 7 ++++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 324e01157b..9c862a2617 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,7 +36,7 @@ huggingface = [
     "tyro",
     "transformers>=4.38.2",
     "datasets>=2.16.0",
-    "sentencepiece",
+    "sentencepiece>=0.2.0",
     "tqdm",
     "psutil",
     "wheel>=0.42.0",
diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 0259ab7889..c0cce75e2f 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -145,7 +145,12 @@ def GemmaModel_fast_forward_inference(
     bsz, q_len, hd = hidden_states.shape
     seq_len = past_key_values[0][0].shape[-2]
     if bsz != 1:
-        attention_mask = _prepare_4d_causal_attention_mask(attention_mask, (bsz, q_len), hidden_states, seq_len,)
+        attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+            attention_mask,
+            (bsz, q_len),
+            hidden_states,
+            seq_len,
+        )
     pass
 
     next_decoder_cache = []