File size: 1,260 Bytes
e95c403 2a17ff8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """Shared loader helper for the examples.
Provides `load_kernel()` which returns the mamba3 kernel package. Uses the
local clone (examples/../build/torch-neuron/) when running from a git clone,
or falls back to the HF Hub via `get_kernel("jburtoft/mamba3-neuron-kernels")`
when installed via the `kernels` library.
"""
import os
import sys
def load_kernel():
"""Return the mamba3 kernel module."""
local_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), "..", "build", "torch-neuron", "__init__.py")
)
if os.path.exists(local_path):
# print(f"[loader] using local clone: {os.path.dirname(local_path)}")
import importlib.util
spec = importlib.util.spec_from_file_location(
"mamba3", local_path,
submodule_search_locations=[os.path.dirname(local_path)],
)
m = importlib.util.module_from_spec(spec)
sys.modules["mamba3"] = m
spec.loader.exec_module(m)
return m
else:
# print("[loader] using HF Hub: jburtoft/mamba3-neuron-kernels")
from kernels import get_kernel
return get_kernel(
"jburtoft/mamba3-neuron-kernels",
revision="v1.0.0",
trust_remote_code=True,
)
|