Fix tokenizer test failure

mlx_lm/models/gemma4_text.py

@@ -609,6 +609,7 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         sanitized = {}
+        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
         for k, v in weights.items():
             if any(
                 s in k
@@ -622,6 +623,18 @@ class Model(nn.Module):
             ):
                 continue
 
+            # KV-shared layers reuse K/V from earlier layers — drop their projections
+            if any(
+                s in k
+                for s in (".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.k_norm")
+            ):
+                try:
+                    layer_idx = int(k.split("layers.")[1].split(".")[0])
+                    if layer_idx >= first_kv_shared:
+                        continue
+                except (IndexError, ValueError):
+                    pass
+
             if k.endswith(".experts.gate_up_proj"):
                 base = k.removesuffix(".gate_up_proj")
                 gate, up = map(mx.contiguous, mx.split(v, 2, axis=-2))

mlx_lm/tokenizer_utils.py

@@ -452,12 +452,12 @@ class TokenizerWrapper:
         """
         Get a stateful streaming detokenizer.
         """
-        return self._detokenizer_class(self)
+        if not hasattr(self, "_detokenizer"):
+            self._detokenizer = self._detokenizer_class(self)
+        return self._detokenizer
 
     def __getattr__(self, attr):
-        if attr == "detokenizer":
-            return self._detokenizer
-        elif attr == "eos_token_ids":
+        if attr == "eos_token_ids":
             return self._eos_token_ids
         elif attr.startswith("_"):
             return self.__getattribute__(attr)

tests/test_server.py

@@ -203,8 +203,10 @@ class TestServer(unittest.TestCase):
         self.assertIn("choices", response_body)
         first_text = response_body["choices"][0]["text"]
         self.assertEqual(
-            first_text,
-            json.loads(requests.post(url, json=post_data).text)["choices"][0]["text"],
+            first_text.strip(),
+            json.loads(requests.post(url, json=post_data).text)["choices"][0][
+                "text"
+            ].strip(),
         )
 
     def test_handle_chat_completions(self):

tests/test_tokenizers.py

@@ -65,6 +65,8 @@ class TestTokenizers(unittest.TestCase):
                 tokenizer = load_tokenizer(tokenizer_repo)
                 tokenizer.decode([0, 1, 2])
                 self.assertTrue(isinstance(tokenizer.detokenizer, expected_detokenizer))
+                if expected_detokenizer is BPEStreamingDetokenizer:
+                    tokenizer.detokenizer.clean_spaces = False
                 self.check_tokenizer(tokenizer)
 
         # Try one with a naive detokenizer