Set multiple eos tokens for VLM (#119)

* Set multiple eos tokens for VLM

* attempt 2

Author: neilmehta24

mlx_engine/generate.py

@@ -267,11 +267,7 @@ def sampler_func_wrapper(*args, **kwargs):
     tokenizer = model_kit.tokenizer
 
     # Set up stop string processor if non-empty stop_strings are provided
-    eos_token_ids = (
-        tokenizer.eos_token_ids
-        if isinstance(tokenizer.eos_token_ids, Iterable)
-        else [tokenizer.eos_token_ids]
-    )
+    eos_token_ids = tokenizer.eos_token_ids
     stop_string_processor = None
     if stop_strings is not None and len(stop_strings) > 0:
         stop_string_processor = StopStringProcessor(stop_strings, tokenizer)

mlx_engine/vision/vision_model_kit.py

@@ -6,6 +6,7 @@
 from .vision_model_wrapper import VisionModelWrapper
 
 import mlx_vlm
+import mlx_lm
 from pathlib import Path
 import mlx.core as mx
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
@@ -100,8 +101,21 @@ def _full_model_init(self):
         else:
             self.model, self.processor, self.model_weights = return_tuple
         self.model = VisionModelWrapper(self.model)
-        self.tokenizer = mlx_vlm.tokenizer_utils.load_tokenizer(self.model_path)
+
+        # Set the eos_token_ids
+        eos_token_ids = []
+        if (eos_tokens := self.config.get("eos_token_ids", None)) is not None:
+            eos_token_ids = list(set(eos_tokens))
+            log_info(f"Setting eos token ids: {eos_token_ids}")
+        elif (eos_tokens := self.config.get("eos_token_id", None)) is not None:
+            eos_token_ids = [eos_tokens]
+
+        # Use the mlx_lm tokenizer since it's more robust
+        self.tokenizer = mlx_lm.tokenizer_utils.load_tokenizer(
+            self.model_path, eos_token_ids=list(eos_token_ids)
+        )
         self.detokenizer = self.tokenizer.detokenizer
+
         self.cache_wrapper = None
         mx.metal.clear_cache()
 

tests/test_vision_models.py

@@ -227,6 +227,7 @@ def test_gemma3_text_only(self):
         prompt = f"{self.text_only_prompt}"
         self.model_helper("mlx-community/gemma-3-4b-it-4bit", prompt, text_only=True)
 
+
 """
 To find the correct prompt format for new models, run this command for your model in the terminal and check the prompt dump:
 python -m mlx_vlm.generate --model ~/.cache/lm-studio/models/mlx-community/MODEL-NAME --max-tokens 100 --temp 0.0 --image http://images.cocodataset.org/val2017/000000039769.jpg --prompt "What do you see?"