From a7cdd56a45f82430f2ae61f4f9aa6b575b74689f Mon Sep 17 00:00:00 2001
From: 35Upesh <35upesh@gmail.com>
Date: Thu, 6 Aug 2026 08:28:59 +0000
Subject: [PATCH] Code changes for google/gemma-4-31b-it model text config

Signed-off-by: 35Upesh <35upesh@gmail.com>
---
 ARCHITECTURE.md                 | 7 ++++---
 README.md                       | 2 +-
 hf_adapters/auto_spyre_model.py | 1 +
 tests/model_registry.py         | 8 ++++++++
 4 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md
index 54fc2351..11c20cad 100644
--- a/ARCHITECTURE.md
+++ b/ARCHITECTURE.md
@@ -31,6 +31,7 @@ which models are supported on Spyre.
 | Gemma 4 12B | gemma4\_unified | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 4 12B Base | gemma4\_unified | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 4 31B | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
+| Gemma 4 31B Instruct | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 3 1B | gemma3\_text | 256 | 128 | Yes | Yes | Yes | Yes |
 | GPT-2 124M | gpt2 | 64 | n/a (no RoPE) | Yes | Yes | Yes | Yes |
 | GPT-Neo 125M | gpt_neo | 64 | n/a (no RoPE) | Yes | Yes | Yes | Yes |
@@ -112,8 +113,8 @@ single-token decode path (seq_len=1), not an adapter issue.
 > adapter or verify a checkpoint, update *only* this file (and the badge
 > counts in README.md, noted below).
 
-**Coverage:** 27 adapters · 46 verified checkpoints · 100+ compatible models.
-The 46 verified rows are 29 generative + 13 embedding + 4 vision-language (see the
+**Coverage:** 27 adapters · 47 verified checkpoints · 100+ compatible models.
+The 47 verified rows are 30 generative + 13 embedding + 4 vision-language (see the
 Verified Checkpoints tables above). `hf_siglip_vision` and `hf_pixtral_vision` are
 vision-tower components used by VLM adapters rather than standalone model adapters.
 Granite Vision 4.1 is verified both as a text backbone (generative) and as a full VLM.
@@ -141,7 +142,7 @@ pattern, norms, and weight layout.
 | hf\_phi3.py | phi3 | 1 | Phi-3 mini 4k/128k, Phi-3 small 8k |
 | hf\_granitemoehybrid.py | granitemoehybrid | 2 | Granite 4.0 Micro |
 | hf\_smollm3.py | smollm3 | 1 | — |
-| hf\_gemma4.py | gemma4\_unified / gemma4 (dense) | 3 | Not E2B/E4B (PLE) or 26B-A4B (MoE). |
+| hf\_gemma4.py | gemma4\_unified / gemma4 (dense) | 4 | Not E2B/E4B (PLE) or 26B-A4B (MoE). |
 | hf\_gemma4\_mm.py | gemma4\_unified (multimodal) | 1 | Gemma 4 31B (dense unified VLM). Not E2B/E4B (PLE) or 26B-A4B (MoE). |
 | hf\_gemma3.py | gemma3\_text / gemma3 (dense) | 2 | Gemma 3 4B/12B/27B (text decoder of the multimodal checkpoints); EmbeddingGemma (bidirectional embedder). Not Gemma 3n (PLE). |
 | hf\_olmo.py | olmo | 1 | OLMo 7B |
diff --git a/README.md b/README.md
index 6e9487c7..4bf6e53e 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ from `transformers`.
 
 ## Supported Models
 
-**30 adapters · 46 verified checkpoints · 100+ compatible models**
+**30 adapters · 47 verified checkpoints · 100+ compatible models**
 
 Coverage spans **generative** (causal-LM), **embedding** (sentence-transformers),
 **vision-language** (image→text), and **speculative-decoding drafter** models — from
diff --git a/hf_adapters/auto_spyre_model.py b/hf_adapters/auto_spyre_model.py
index dbaa8650..abe37e9e 100644
--- a/hf_adapters/auto_spyre_model.py
+++ b/hf_adapters/auto_spyre_model.py
@@ -187,6 +187,7 @@
     "google/gemma-4-12b": torch.bfloat16,
     "google/gemma-4-12B-it": torch.bfloat16,
     "google/gemma-4-31b": torch.bfloat16,
+    "google/gemma-4-31b-it": torch.bfloat16,
     "ibm-granite/granite-4.0-1b-base": torch.float32,
     "ibm-granite/granite-4.0-1b": torch.float32,
 }
diff --git a/tests/model_registry.py b/tests/model_registry.py
index cd1335dd..3255df5d 100644
--- a/tests/model_registry.py
+++ b/tests/model_registry.py
@@ -240,6 +240,14 @@ def _include_gated() -> bool:
         "dtype": "bfloat16",
         "is_gated": True,
     },
+    "gemma4_31b_it": {
+        "name": "Gemma 4 31B Instruct",
+        "path": "google/gemma-4-31b-it",
+        "adapter": "hf_gemma4.py",
+        "size": "31b",
+        "dtype": "bfloat16",
+        "is_gated": True,
+    },
     # DSpark speculative-decoding drafters (block proposers). kind="dspark_draft"
     # keeps them out of the generate-based causal-LM harnesses (see CAUSAL_PATHS);
     # tests/spyre/test_dspark_draft_spyre.py exercises the block-propose path.