feat(lib): add _assets module for bundled-weight resolution
Provide bundled_path() and resolve() helpers that locate ONNX weights and JSON metadata via importlib.resources, falling back to a user-supplied directory. Missing assets raise ModelLoadError. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
52
src/ai_mouse/_assets.py
Normal file
52
src/ai_mouse/_assets.py
Normal file
@@ -0,0 +1,52 @@
|
|||||||
|
"""Asset path resolution for bundled ONNX weights and JSON metadata.
|
||||||
|
|
||||||
|
Uses :mod:`importlib.resources` to locate files inside the installed
|
||||||
|
package, falling back to a user-supplied directory if provided.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from importlib.resources import as_file, files
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from ai_mouse.errors import ModelLoadError
|
||||||
|
|
||||||
|
_PACKAGE_ASSETS = "ai_mouse.assets"
|
||||||
|
|
||||||
|
|
||||||
|
def bundled_path(name: str) -> Path:
|
||||||
|
"""Return a filesystem path to a bundled asset.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: filename inside the assets/ directory.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A concrete :class:`pathlib.Path`. For non-zip wheels (the common
|
||||||
|
case) the path points directly into the installed package; for
|
||||||
|
zipapp installations :func:`importlib.resources.as_file`
|
||||||
|
materialises a temp file.
|
||||||
|
"""
|
||||||
|
ref = files(_PACKAGE_ASSETS) / name
|
||||||
|
with as_file(ref) as p:
|
||||||
|
return Path(p)
|
||||||
|
|
||||||
|
|
||||||
|
def resolve(model_path: Path | None, filename: str) -> Path:
|
||||||
|
"""Locate an asset given an optional user-supplied directory.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_path: user-supplied directory, or None to use bundled assets.
|
||||||
|
filename: file to locate inside the directory.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Absolute path to the asset.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ModelLoadError: if the file does not exist.
|
||||||
|
"""
|
||||||
|
if model_path is None:
|
||||||
|
p = bundled_path(filename)
|
||||||
|
else:
|
||||||
|
p = Path(model_path) / filename
|
||||||
|
if not p.exists():
|
||||||
|
raise ModelLoadError(f"Required asset missing: {p}")
|
||||||
|
return p
|
||||||
34
tests/unit/test_assets.py
Normal file
34
tests/unit/test_assets.py
Normal file
@@ -0,0 +1,34 @@
|
|||||||
|
"""Test the asset path resolver."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from ai_mouse import _assets
|
||||||
|
from ai_mouse.errors import ModelLoadError
|
||||||
|
|
||||||
|
|
||||||
|
def test_bundled_flow_model_exists() -> None:
|
||||||
|
p = _assets.bundled_path("flow_model.onnx")
|
||||||
|
assert p.exists()
|
||||||
|
assert p.suffix == ".onnx"
|
||||||
|
|
||||||
|
|
||||||
|
def test_bundled_train_config_loadable() -> None:
|
||||||
|
p = _assets.bundled_path("train_config.json")
|
||||||
|
cfg = json.loads(p.read_text())
|
||||||
|
assert "seq_len" in cfg
|
||||||
|
assert "d_model" in cfg
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_with_custom_dir(tmp_path: Path) -> None:
|
||||||
|
(tmp_path / "flow_model.onnx").write_bytes(b"x")
|
||||||
|
p = _assets.resolve(tmp_path, "flow_model.onnx")
|
||||||
|
assert p == tmp_path / "flow_model.onnx"
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None:
|
||||||
|
with pytest.raises(ModelLoadError, match="missing"):
|
||||||
|
_assets.resolve(tmp_path, "nonexistent.onnx")
|
||||||
Reference in New Issue
Block a user