Fix server 404 on short prompts: clamp negative start in think-token search (ml-explore#1327)

mlx_lm.server computes start = len(prompt) - 11 and passes it to
rfind_think_start; for prompts under 11 tokens start is negative, so _find
indexes past the token list and raises IndexError (surfaced as HTTP 404) on
common one-word messages like "hi". Clamp start to 0 in _find; add a
regression test.

Author: devYRPauli

mlx_lm/tokenizer_utils.py

@@ -354,8 +354,9 @@ def add_eos_token(self, token: str):
 
         self._eos_token_ids.add(token_id)
 
-    def _find(self, tokens, sequence, start=None, end=None, reverse=False):
-        start = start or 0
+    @staticmethod
+    def _find(tokens, sequence, start=None, end=None, reverse=False):
+        start = max(start or 0, 0)
         end = end or len(tokens)
         outer_loop = (
             range(end - len(sequence), start - 1, -1)

tests/test_tokenizers.py

@@ -9,6 +9,7 @@
     BPEStreamingDetokenizer,
     NaiveStreamingDetokenizer,
     SPMStreamingDetokenizer,
+    TokenizerWrapper,
 )
 from mlx_lm.utils import load_tokenizer
 
@@ -109,6 +110,19 @@ def test_thinking(self):
         self.assertIsNone(tokenizer.think_start_id)
         self.assertIsNone(tokenizer.think_end_id)
 
+    def test_find_token(self):
+        # Check that _find returns a valid index when
+        # searching for a think token in short system prompts
+        HI, THINK_START, THINK_END = 200, 100, 101
+        find = TokenizerWrapper._find
+        prompt = [HI]
+        start = len(prompt) - 11
+        self.assertEqual(find(prompt, [THINK_START], start=start), -1)
+        self.assertEqual(find(prompt, [THINK_START], start=start, reverse=True), -1)
+        prompt = [HI, THINK_START, THINK_END, THINK_START]
+        self.assertEqual(find(prompt, [THINK_START], start=0), 1)
+        self.assertEqual(find(prompt, [THINK_START], start=0, reverse=True), 3)
+
 
 if __name__ == "__main__":
     unittest.main()