ommlds 0.0.0.dev466__py3-none-any.whl → 0.0.0.dev512__py3-none-any.whl

ommlds/backends/mlx/limits.py

@@ -19,9 +19,13 @@ import contextlib
 import sys
 import typing as ta
 
-import mlx.core as mx
-import mlx.utils
-from mlx import nn
+from omlish import lang
+
+
+with lang.auto_proxy_import(globals()):
+    import mlx.core as mx
+    import mlx.nn as mlx_nn
+    import mlx.utils as mlx_utils
 
 
 ##
@@ -29,8 +33,8 @@ from mlx import nn
 
 @contextlib.contextmanager
 def wired_limit_context(
-        model: nn.Module,
-        streams: ta.Iterable[mx.Stream] | None = None,
+        model: 'mlx_nn.Module',
+        streams: ta.Iterable['mx.Stream'] | None = None,
 ) -> ta.Generator[None]:
     """
     A context manager to temporarily change the wired limit.
@@ -43,7 +47,7 @@ def wired_limit_context(
         yield
         return
 
-    model_bytes = mlx.utils.tree_reduce(
+    model_bytes = mlx_utils.tree_reduce(
         lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc,
         model,
         0,

ommlds/backends/mlx/loading.py

@@ -1,10 +1,8 @@
+# ruff: noqa: TC002
 import dataclasses as dc
 import pathlib
 import typing as ta
 
-import mlx_lm.utils
-from mlx import nn
-
 from omlish import check
 from omlish import lang
 
@@ -12,6 +10,11 @@ from .tokenization import Tokenization
 from .tokenization import load_tokenization
 
 
+with lang.auto_proxy_import(globals()):
+    import mlx.nn as mlx_nn
+    import mlx_lm.utils
+
+
 ##
 
 
@@ -76,7 +79,7 @@ def get_model_path(
 class LoadedModel:
     path: pathlib.Path
 
-    model: nn.Module
+    model: 'mlx_nn.Module'
     config: dict
 
     #

ommlds/backends/ollama/__init__.py

@@ -0,0 +1,7 @@
+from omlish import dataclasses as _dc  # noqa
+
+
+_dc.init_package(
+    globals(),
+    codegen=True,
+)