| """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) |
|
|
