File size: 5,852 Bytes
8d20d7d | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Explicit local inference selection without importing optional ML dependencies."""
import importlib
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from typing import Protocol, cast
from .suggestion import SuggestionResult, suggest
_OPTIONAL_MODULES = frozenset({"numpy", "packaging", "safetensors", "torch"})
class InferenceError(ValueError):
"""An explicit inference backend cannot be configured or loaded safely."""
class InferenceBackend(StrEnum):
"""Available instruction interpreters; baseline remains the default."""
BASELINE = "baseline"
LEARNED = "learned"
def _validate_device(device: object) -> None:
"""Reject invalid or unhashable device values at the runtime boundary."""
if not isinstance(device, str) or device not in {"cpu", "mps"}:
raise InferenceError("Device must be explicitly 'cpu' or 'mps'.")
@dataclass(frozen=True, slots=True)
class InferenceConfig:
"""Explicit backend, acknowledgement, checkpoint, and device selection."""
backend: InferenceBackend = InferenceBackend.BASELINE
allow_experimental: bool = False
config_path: Path | None = None
weights_path: Path | None = None
device: str = "cpu"
def __post_init__(self) -> None:
"""Reject ignored or incomplete backend options."""
if type(self.backend) is not InferenceBackend:
raise InferenceError("Backend must be an InferenceBackend enum value.")
if type(self.allow_experimental) is not bool:
raise InferenceError("Experimental acknowledgement must be an exact boolean.")
if self.config_path is not None and not isinstance(self.config_path, Path):
raise InferenceError("Config checkpoint must be a Path or None.")
if self.weights_path is not None and not isinstance(self.weights_path, Path):
raise InferenceError("Weights checkpoint must be a Path or None.")
_validate_device(self.device)
learned_options = (
self.allow_experimental or self.config_path is not None or self.weights_path is not None
)
if self.backend == InferenceBackend.BASELINE:
if learned_options or self.device != "cpu":
raise InferenceError(
"Checkpoint, device, and opt-in options require learned backend."
)
return
if not self.allow_experimental:
raise InferenceError("Learned inference requires allow_experimental=True.")
if self.config_path is None or self.weights_path is None:
raise InferenceError("Learned inference requires explicit config and weights paths.")
class LearnedRuntime(Protocol):
"""Minimal typed boundary implemented by the optional learned runtime."""
@property
def device(self) -> str:
"""Return the device actually selected without fallback."""
...
def __call__(self, source: str, instruction: str) -> SuggestionResult:
"""Support the existing evaluation predictor contract."""
...
def suggest(self, source: str, instruction: str) -> SuggestionResult:
"""Return one immutable core suggestion result."""
...
class _RuntimeType(Protocol):
def load(self, config_path: Path, weights_path: Path, device: str) -> LearnedRuntime:
"""Load one runtime through the hardened optional implementation."""
...
def _load_learned_runtime(config_path: Path, weights_path: Path, device: str) -> LearnedRuntime:
module = importlib.import_module(".ml.runtime", __package__)
runtime_type = cast("_RuntimeType", module.__dict__["LoadedLearnedRuntime"])
return runtime_type.load(config_path, weights_path, device)
def _optional_dependency_error(error: ModuleNotFoundError) -> InferenceError | None:
name = error.name or ""
top_level = name.partition(".")[0]
if top_level not in _OPTIONAL_MODULES:
return None
return InferenceError(
"Learned inference dependencies are unavailable; install the local package with "
"the 'yaml-edit[ml]' extra."
)
class InferenceEngine:
"""Reusable explicit inference adapter that loads at most one learned model."""
def __init__(self, config: InferenceConfig | None = None) -> None:
"""Validate configuration and eagerly load one requested local checkpoint."""
self.config = config if config is not None else InferenceConfig()
self._runtime: LearnedRuntime | None = None
if self.config.backend == InferenceBackend.LEARNED:
self._runtime = self._load_runtime()
def _load_runtime(self) -> LearnedRuntime:
config_path = self.config.config_path
weights_path = self.config.weights_path
if config_path is None or weights_path is None:
raise InferenceError("Learned inference requires explicit checkpoint paths.")
try:
return _load_learned_runtime(config_path, weights_path, self.config.device)
except ModuleNotFoundError as error:
classified = _optional_dependency_error(error)
if classified is None:
raise
raise classified from error
except (OSError, ValueError) as error:
raise InferenceError(
f"Learned inference could not load the checkpoint: {error}"
) from error
@property
def runtime(self) -> LearnedRuntime | None:
"""Expose the one explicitly loaded runtime for evaluation and measurement."""
return self._runtime
def suggest(self, source: str, instruction: str) -> SuggestionResult:
"""Interpret without applying or writing the source document."""
if self._runtime is None:
return suggest(source, instruction)
return self._runtime.suggest(source, instruction)
|