mlxsmith 0.1.0__py3-none-any.whl

mlxsmith/__init__.py

@@ -0,0 +1,2 @@
+__all__ = ['__version__']
+__version__ = '0.1.0'

mlxsmith/accel/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+from .none import NoneBackend
+from .zmlx_backend import ZMLXBackend
+
+def get_backend(name: str):
+    if name == "none":
+        return NoneBackend()
+    if name == "zmlx":
+        return ZMLXBackend()
+    raise ValueError(f"Unknown accel backend: {name}")

mlxsmith/accel/base.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Protocol
+
+@dataclass
+class AccelStats:
+    backend: str
+    compiled: int = 0
+    cache_hits: int = 0
+    notes: Dict[str, Any] | None = None
+
+class AccelBackend(Protocol):
+    name: str
+    def patch(self) -> None: ...
+    def warmup(self, model: Any, example_batch: Any) -> Dict[str, Any]: ...
+    def stats(self) -> AccelStats: ...

mlxsmith/accel/none.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+from .base import AccelStats
+
+class NoneBackend:
+    name = "none"
+    def patch(self) -> None:
+        return
+    def warmup(self, model: Any, example_batch: Any) -> Dict[str, Any]:
+        return {"warmup": "skipped"}
+    def stats(self) -> AccelStats:
+        return AccelStats(backend="none")

mlxsmith/accel/zmlx_backend.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+from .base import AccelStats
+
+class ZMLXBackend:
+    name = "zmlx"
+
+    def __init__(self):
+        self._available = False
+        self._notes = {}
+        try:
+            import zmlx  # type: ignore
+            self._available = True
+            self._notes["zmlx_version"] = getattr(zmlx, "__version__", None)
+        except Exception as e:
+            self._available = False
+            self._notes["error"] = f"{type(e).__name__}: {e}"
+
+    def patch(self) -> None:
+        if not self._available:
+            # soft fail; caller should report status
+            return
+        # ZMLX can patch ops/modules. We keep this intentionally minimal and safe.
+        try:
+            import zmlx  # type: ignore
+            # If ZMLX provides a global patch hook, call it; otherwise, no-op.
+            patch_fn = getattr(zmlx, "patch", None)
+            if callable(patch_fn):
+                patch_fn()
+                self._notes["patched"] = True
+            else:
+                self._notes["patched"] = False
+                self._notes["hint"] = "No zmlx.patch() found; implement patch hook or integrate per-module."
+        except Exception as e:
+            self._notes["patched_error"] = f"{type(e).__name__}: {e}"
+
+    def warmup(self, model: Any, example_batch: Any) -> Dict[str, Any]:
+        return {"warmup": "not_implemented", "notes": self._notes}
+
+    def stats(self) -> AccelStats:
+        return AccelStats(backend="zmlx", notes=self._notes)

mlxsmith/adapters.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Iterable, Optional
+
+from .util import ensure_dir, copytree, now_ts
+
+
+def merge_adapters(
+    base_model: str,
+    adapters: Iterable[Path],
+    out_dir: Path,
+    *,
+    weights: Optional[list[float]] = None,
+) -> Path:
+    adapter_list = [Path(a) for a in adapters]
+    if not adapter_list:
+        raise RuntimeError("No adapters provided")
+    ensure_dir(out_dir)
+
+    try:
+        from mlx_lm.tuner import utils as tuner_utils  # type: ignore
+
+        if hasattr(tuner_utils, "merge_adapters"):
+            tuner_utils.merge_adapters(
+                base_model,
+                [str(p) for p in adapter_list],
+                str(out_dir),
+                weights=weights,
+            )
+            return out_dir
+    except Exception:
+        pass
+
+    # Fallback: copy the first adapter and record metadata.
+    copytree(adapter_list[0], out_dir)
+    meta = {
+        "base_model": base_model,
+        "merged_from": [str(p) for p in adapter_list],
+        "weights": weights,
+        "merged_at": now_ts(),
+        "note": "merge_adapters fallback: copied first adapter (mlx_lm merge unavailable)",
+    }
+    (out_dir / "adapter_merge.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
+    return out_dir

mlxsmith/api/__init__.py

@@ -0,0 +1,48 @@
+"""MLXSmith API handlers and schemas."""
+
+from .schemas import (
+    ChatMessage,
+    ChatRequest,
+    ChatResponse,
+    ChatCompletionChunk,
+    RolloutRequest,
+    RolloutResponse,
+    AdapterReloadRequest,
+    AdapterReloadResponse,
+    RLMState,
+    RLMHistoryEntry,
+    ModelInfo,
+    ModelsListResponse,
+    ModelPullRequest,
+    ModelPullResponse,
+    HFTokenRequest,
+    HFTokenResponse,
+    HealthResponse,
+    ErrorResponse,
+)
+from .handlers import create_router, InternalAuthMiddleware
+
+__all__ = [
+    # Schemas
+    "ChatMessage",
+    "ChatRequest",
+    "ChatResponse",
+    "ChatCompletionChunk",
+    "RolloutRequest",
+    "RolloutResponse",
+    "AdapterReloadRequest",
+    "AdapterReloadResponse",
+    "RLMState",
+    "RLMHistoryEntry",
+    "ModelInfo",
+    "ModelsListResponse",
+    "ModelPullRequest",
+    "ModelPullResponse",
+    "HFTokenRequest",
+    "HFTokenResponse",
+    "HealthResponse",
+    "ErrorResponse",
+    # Handlers
+    "create_router",
+    "InternalAuthMiddleware",
+]