linkml-store 0.3.0__py3-none-any.whl

linkml_store/inference/inference_engine.py

@@ -0,0 +1,209 @@
+import logging
+import random
+from abc import ABC
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any, Optional, TextIO, Tuple, Union
+
+import pandas as pd
+from pydantic import BaseModel, ConfigDict
+
+from linkml_store.api.collection import OBJECT, Collection
+from linkml_store.inference.inference_config import Inference, InferenceConfig
+from linkml_store.utils.pandas_utils import nested_objects_to_dataframe
+
+logger = logging.getLogger(__name__)
+
+
+class ModelSerialization(str, Enum):
+    """
+    Enum for model serialization types.
+    """
+
+    PICKLE = "pickle"
+    ONNX = "onnx"
+    PMML = "pmml"
+    PFA = "pfa"
+    JOBLIB = "joblib"
+    PNG = "png"
+    LINKML_EXPRESSION = "linkml_expression"
+    RULE_BASED = "rulebased"
+    RAG_INDEX = "rag_index"
+
+    @classmethod
+    def from_filepath(cls, file_path: str) -> Optional["ModelSerialization"]:
+        """
+        Get the serialization type from the file path.
+
+        >>> ModelSerialization.from_filepath("model.onnx")
+        <ModelSerialization.ONNX: 'onnx'>
+        >>> ModelSerialization.from_filepath("model.pkl")
+        <ModelSerialization.PICKLE: 'pickle'>
+        >>> assert ModelSerialization.from_filepath("poor_file_name") is None
+
+        :param file_path:
+        :return:
+        """
+        toks = file_path.split(".")
+        suffix = toks[-1]
+        if len(toks) > 2:
+            if suffix == "yaml" and toks[-2] == "rulebased":
+                return cls.RULE_BASED
+        # Generate mapping dynamically
+        extension_mapping = {v.lower(): v for v in cls}
+        # Add special cases
+        extension_mapping["pkl"] = cls.PICKLE
+        extension_mapping["py"] = cls.LINKML_EXPRESSION
+        return extension_mapping.get(suffix, None)
+
+
+class CollectionSlice(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+
+    name: Optional[str] = None
+    base_collection: Optional[Collection] = None
+    # _dataframe: Optional[pd.DataFrame] = None
+    # slice: Tuple[Optional[int], Optional[int]] = Field(default=(None, None))
+    indices: Optional[Tuple[int, ...]] = None
+    _collection: Optional[Collection] = None
+    where: Any = None
+
+    @property
+    def collection(self) -> Collection:
+        if not self._collection and not self.indices:
+            return self.base_collection
+        if not self._collection:
+            rows = self.base_collection.rows
+            subset = [rows[i] for i in self.indices]
+            db = self.base_collection.parent
+            subset_name = self.slice_alias
+            subset_collection = db.get_collection(subset_name, create_if_not_exists=True)
+            # ensure the collection has the same schema type as the base collection;
+            # this ensures that column/attribute types are preserved
+            subset_collection.metadata.type = self.base_collection.target_class_name
+            subset_collection.delete_where({})
+            subset_collection.insert(subset)
+            self._collection = subset_collection
+        return self._collection
+
+    @property
+    def slice_alias(self) -> str:
+        return f"{self.base_collection.alias}__rag_{self.name}"
+
+    def as_dataframe(self, flattened=False) -> pd.DataFrame:
+        """
+        Return the slice of the collection as a dataframe.
+
+        :param flattened: flattned nested objects to give keys like foo.bar
+        :return:
+        """
+        rs = self.collection.find({}, limit=-1)
+        if flattened:
+            return nested_objects_to_dataframe(rs.rows)
+        else:
+            return rs.rows_dataframe
+
+
+@dataclass
+class InferenceEngine(ABC):
+    """
+    Base class for all inference engine.
+
+    An InferenceEngine is capable of deriving inferences from input objects and a collection.
+    """
+
+    predictor_type: Optional[str] = None
+    config: Optional[InferenceConfig] = None
+
+    training_data: Optional[CollectionSlice] = None
+    testing_data: Optional[CollectionSlice] = None
+
+    def load_and_split_data(self, collection: Collection, split: Optional[Tuple[float, float]] = None, randomize=True):
+        """
+        Load the data and split it into training and testing sets.
+
+        :param collection:
+        :param split: Tuple of training and testing split ratios.
+        :param randomize:
+        :return:
+        """
+        local_random = random.Random(self.config.random_seed) if self.config.random_seed else random.Random()
+        split = split or self.config.train_test_split
+        if not split:
+            split = (0.7, 0.3)
+        if split[0] == 1.0:
+            self.training_data = CollectionSlice(name="train", base_collection=collection, indices=None)
+            self.testing_data = None
+            return
+        logger.info(f"Loading and splitting data {split} from collection {collection.alias}")
+        size = collection.size()
+        indices = range(size)
+        if randomize:
+            train_indices = local_random.sample(indices, int(size * split[0]))
+            test_indices = set(indices) - set(train_indices)
+        else:
+            train_indices = indices[: int(size * split[0])]
+            test_indices = indices[int(size * split[0]) :]
+        self.training_data = CollectionSlice(name="train", base_collection=collection, indices=train_indices)
+        self.testing_data = CollectionSlice(name="test", base_collection=collection, indices=test_indices)
+
+    def initialize_model(self, **kwargs):
+        """
+        Initialize the model.
+
+        :param kwargs:
+        :return:
+        """
+        raise NotImplementedError("Initialize model method must be implemented by subclass")
+
+    def export_model(
+        self, output: Optional[Union[str, Path, TextIO]], model_serialization: ModelSerialization = None, **kwargs
+    ):
+        """
+        Export the model to the given output.
+
+        :param model_serialization:
+        :param output:
+        :param kwargs:
+        :return:
+        """
+        raise NotImplementedError("Export model method must be implemented by subclass")
+
+    def import_model_from(self, inference_engine: "InferenceEngine", **kwargs):
+        """
+        Import the model from the given inference engine.
+
+        :param inference_engine:
+        :param kwargs:
+        :return:
+        """
+        raise NotImplementedError("Import model method must be implemented by subclass")
+
+    def save_model(self, output: Union[str, Path]) -> None:
+        """
+        Save the model to the given output.
+
+        :param output:
+        :return:
+        """
+        raise NotImplementedError("Save model method must be implemented by subclass")
+
+    @classmethod
+    def load_model(cls, file_path: Union[str, Path]) -> "InferenceEngine":
+        """
+        Load the model from the given file path.
+
+        :param file_path:
+        :return:
+        """
+        raise NotImplementedError("Load model method must be implemented by subclass")
+
+    def derive(self, object: OBJECT) -> Optional[Inference]:
+        """
+        Derive the prediction for the given object.
+
+        :param object:
+        :return:
+        """
+        raise NotImplementedError("Predict method must be implemented by subclass")

linkml_store/inference/inference_engine_registry.py

@@ -0,0 +1,74 @@
+import importlib
+import inspect
+import os
+from typing import Dict, Type
+
+from linkml_store.inference.inference_config import InferenceConfig
+from linkml_store.inference.inference_engine import InferenceEngine
+from linkml_store.utils.object_utils import object_path_update
+
+
+class InferenceEngineRegistry:
+    def __init__(self):
+        self.engines: Dict[str, Type[InferenceEngine]] = {}
+
+    def register(self, name: str, engine_class: Type[InferenceEngine]):
+        self.engines[name] = engine_class
+
+    def get_engine_class(self, name: str) -> Type[InferenceEngine]:
+        if name not in self.engines:
+            raise ValueError(f"Unknown inference engine type: {name}" f"Known engines: {list(self.engines.keys())}")
+        return self.engines[name]
+
+    def create_engine(self, engine_type: str, config: InferenceConfig = None, **kwargs) -> InferenceEngine:
+        kwargs = {k: v for k, v in kwargs.items() if v is not None}
+        if ":" in engine_type:
+            engine_type, conf_args = engine_type.split(":", 1)
+            if config is None:
+                config = InferenceConfig()
+            for arg in conf_args.split(","):
+                k, v = arg.split("=")
+                config = object_path_update(config, k, v)
+
+        engine_class = self.get_engine_class(engine_type)
+        kwargs["predictor_type"] = engine_type
+        return engine_class(config=config, **kwargs)
+
+    @classmethod
+    def load_engines(cls, package_path: str):
+        registry = cls()
+        package_dir = os.path.dirname(importlib.import_module(package_path).__file__)
+        for filename in os.listdir(package_dir):
+            if filename.endswith(".py") and not filename.startswith("__"):
+                module_name = f"{package_path}.{filename[:-3]}"
+                try:
+                    module = importlib.import_module(module_name)
+                    for name, obj in inspect.getmembers(module):
+                        if inspect.isclass(obj) and issubclass(obj, InferenceEngine) and obj != InferenceEngine:
+                            engine_name = name.lower().replace("inferenceengine", "")
+                            registry.register(engine_name, obj)
+                except ImportError as e:
+                    print(f"Error importing {module_name}: {e}")
+        return registry
+
+
+# Initialize the registry
+registry = InferenceEngineRegistry.load_engines("linkml_store.inference.implementations")
+
+
+# Function to get an inference engine (can be used as before)
+def get_inference_engine(engine_type: str, config: InferenceConfig = None, **kwargs) -> InferenceEngine:
+    """
+    Get an inference engine.
+
+    >>> from linkml_store.inference import get_inference_engine
+    >>> ie = get_inference_engine('sklearn')
+    >>> type(ie)
+    <class 'linkml_store.inference.implementations.sklearn_inference_engine.SklearnInferenceEngine'>
+
+    :param engine_type:
+    :param config:
+    :param kwargs:
+    :return:
+    """
+    return registry.create_engine(engine_type, config, **kwargs)

linkml_store/plotting/__init__.py

@@ -0,0 +1,5 @@
+"""
+Visualization and plotting functions for LinkML data.
+"""
+
+__version__ = "0.1.0"