From d3bf847e6fd506b07370813c1fa3b666b4f8f576 Mon Sep 17 00:00:00 2001
From: tnadav <tnadav@gmail.com>
Date: Wed, 5 Nov 2025 21:25:00 +0200
Subject: [PATCH] Make mlx-lm more type-checker friendly (#573)

* Fix type annotation for `load` parameter

* Add type annotations to all `load` parameters

* Avoid using mutable types for `load` default parameters

* Add return type annotation to `load_tokenizer`

* Export public module attributes
---
 mlx_lm/__init__.py        |  9 +++++++++
 mlx_lm/tokenizer_utils.py | 12 ++++++++----
 mlx_lm/utils.py           | 11 ++++++-----
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/mlx_lm/__init__.py b/mlx_lm/__init__.py
index 7180a42..7fb4194 100644
--- a/mlx_lm/__init__.py
+++ b/mlx_lm/__init__.py
@@ -9,3 +9,12 @@ os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 from .convert import convert
 from .generate import batch_generate, generate, stream_generate
 from .utils import load
+
+__all__ = [
+    "__version__",
+    "convert",
+    "batch_generate",
+    "generate",
+    "stream_generate",
+    "load",
+]
diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py
index f6bcd98..c612c4b 100644
--- a/mlx_lm/tokenizer_utils.py
+++ b/mlx_lm/tokenizer_utils.py
@@ -1,7 +1,7 @@
 import json
 from functools import partial
 from json import JSONDecodeError
-from typing import List
+from typing import Any, Dict, List, Optional
 
 from transformers import AutoTokenizer, PreTrainedTokenizerFast
 
@@ -424,8 +424,11 @@ def _is_bpe_decoder(decoder):
 
 
 def load_tokenizer(
-    model_path, tokenizer_config_extra={}, return_tokenizer=True, eos_token_ids=None
-):
+    model_path,
+    tokenizer_config_extra: Optional[Dict[str, Any]] = None,
+    return_tokenizer=True,
+    eos_token_ids=None,
+) -> TokenizerWrapper:
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -454,8 +457,9 @@ def load_tokenizer(
         eos_token_ids = [eos_token_ids]
 
     if return_tokenizer:
+        kwargs = tokenizer_config_extra or {}
         return TokenizerWrapper(
-            AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
+            AutoTokenizer.from_pretrained(model_path, **kwargs),
             detokenizer_class,
             eos_token_ids=eos_token_ids,
         )
diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py
index a3c209c..cac80cb 100644
--- a/mlx_lm/utils.py
+++ b/mlx_lm/utils.py
@@ -148,7 +148,7 @@ def load_model(
     model_path: Path,
     lazy: bool = False,
     strict: bool = True,
-    model_config: dict = {},
+    model_config: Optional[Dict[str, Any]] = None,
     get_model_classes: Callable[[dict], Tuple[Type[nn.Module], Type]] = _get_classes,
 ) -> Tuple[nn.Module, dict]:
     """
@@ -175,7 +175,8 @@ def load_model(
         ValueError: If the model class or args class are not found or cannot be instantiated.
     """
     config = load_config(model_path)
-    config.update(model_config)
+    if model_config is not None:
+        config.update(model_config)
 
     weight_files = glob.glob(str(model_path / "model*.safetensors"))
 
@@ -245,12 +246,12 @@ def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
 
 def load(
     path_or_hf_repo: str,
-    tokenizer_config={},
-    model_config={},
+    tokenizer_config: Optional[Dict[str, Any]] = None,
+    model_config: Optional[Dict[str, Any]] = None,
     adapter_path: Optional[str] = None,
     lazy: bool = False,
     return_config: bool = False,
-    revision: str = None,
+    revision: Optional[str] = None,
 ) -> Union[
     Tuple[nn.Module, TokenizerWrapper],
     Tuple[nn.Module, TokenizerWrapper, Dict[str, Any]],