Files
sillyhome-next-dev/app/ml/predictor.py

50 lines
1.9 KiB
Python

from __future__ import annotations
import logging
from typing import Sequence
from app.ml.feature_store import FeatureVector
from app.ml.registry.model_registry import ModelRegistry
from app.ml.training import TrainedArtifact, TrainingPipeline
logger = logging.getLogger(__name__)
class Predictor:
def __init__(
self,
pipeline: TrainingPipeline | None = None,
registry: ModelRegistry | None = None,
) -> None:
if isinstance(pipeline, ModelRegistry) and registry is None:
registry = pipeline
pipeline = None
if pipeline is None and registry is None:
raise ValueError("Predictor erfordert TrainingPipeline oder ModelRegistry.")
self._pipeline = pipeline
self._registry = registry
def predict(self, artifact_id: str, entity: FeatureVector) -> str:
artifact = self._get_artifact(artifact_id)
if entity.sensor_id not in artifact.supported_sensors:
raise ValueError(
f"Sensor '{entity.sensor_id}' wird vom Modell '{artifact_id}' nicht unterstützt."
)
return f"{artifact_id}:{entity.sensor_id}:{entity.values}"
def predict_batch(self, artifact_id: str, entities: Sequence[FeatureVector]) -> list[str]:
return [self.predict(artifact_id, entity) for entity in entities]
@staticmethod
def default_artifact(pipeline: TrainingPipeline) -> TrainedArtifact:
artifacts = list(pipeline._artifacts)
if not artifacts:
raise ValueError("Kein trainiertes Modell gefunden.")
return pipeline.export(artifacts[-1])
def _get_artifact(self, artifact_id: str) -> TrainedArtifact:
if self._registry is not None:
return self._registry.load_artifact(artifact_id)
if self._pipeline is not None:
return self._pipeline.export(artifact_id)
raise RuntimeError("Predictor nicht initialisiert.")