linkml-store 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

linkml_store/inference/implementations/sklearn_inference_engine.py

@@ -0,0 +1,308 @@
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, ClassVar, Dict, List, Optional, TextIO, Type, Union
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import cross_val_score
+from sklearn.preprocessing import LabelEncoder, MultiLabelBinarizer, OneHotEncoder
+from sklearn.tree import DecisionTreeClassifier
+
+from linkml_store.api.collection import OBJECT
+from linkml_store.inference.implementations.rule_based_inference_engine import RuleBasedInferenceEngine
+from linkml_store.inference.inference_config import Inference, InferenceConfig
+from linkml_store.inference.inference_engine import InferenceEngine, ModelSerialization
+from linkml_store.utils.sklearn_utils import tree_to_nested_expression, visualize_decision_tree
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SklearnInferenceEngine(InferenceEngine):
+    config: InferenceConfig
+    classifier: Any = None
+    encoders: Dict[str, Any] = field(default_factory=dict)
+    transformed_features: List[str] = field(default_factory=list)
+    transformed_targets: List[str] = field(default_factory=list)
+    skip_features: List[str] = field(default_factory=list)
+    categorical_encoder_class: Optional[Type[Union[OneHotEncoder, MultiLabelBinarizer]]] = None
+    maximum_proportion_distinct_features: float = 0.2
+    confidence: float = 0.0
+
+    strict: bool = False
+
+    PERSIST_COLS: ClassVar = [
+        "config",
+        "classifier",
+        "encoders",
+        "transformed_features",
+        "transformed_targets",
+        "skip_features",
+        "confidence",
+    ]
+
+    def _get_encoder(self, v: Union[List[Any], Any]) -> Any:
+        if isinstance(v, list):
+            if all(isinstance(x, list) for x in v):
+                return MultiLabelBinarizer()
+            elif all(isinstance(x, str) for x in v):
+                return OneHotEncoder(sparse_output=False, handle_unknown="ignore")
+            elif all(isinstance(x, (int, float)) for x in v):
+                return None
+            else:
+                raise ValueError("Mixed data types in the list are not supported")
+        else:
+            if hasattr(v, "dtype"):
+                if v.dtype == "object" or v.dtype.name == "category":
+                    if isinstance(v.iloc[0], list):
+                        return MultiLabelBinarizer()
+                    elif self.categorical_encoder_class:
+                        return self.categorical_encoder_class(handle_unknown="ignore")
+                    else:
+                        return OneHotEncoder(sparse_output=False, handle_unknown="ignore")
+                elif v.dtype.kind in "biufc":
+                    return None
+        raise ValueError("Unable to determine appropriate encoder for the input data")
+
+    def _is_complex_column(self, column: pd.Series) -> bool:
+        """Check if the column contains complex data types like lists or dicts."""
+        # MV_TYPE = (list, dict)
+        MV_TYPE = (list,)
+        return (column.dtype == "object" or column.dtype == "category") and any(
+            isinstance(x, MV_TYPE) for x in column.dropna()
+        )
+
+    def _get_unique_values(self, column: pd.Series) -> set:
+        """Get unique values from a column, handling list-type data."""
+        if self._is_complex_column(column):
+            # For columns with lists, flatten the lists and get unique values
+            return set(
+                item for sublist in column.dropna() for item in (sublist if isinstance(sublist, list) else [sublist])
+            )
+        else:
+            return set(column.unique())
+
+    def initialize_model(self, **kwargs):
+        logger.info(f"Initializing model with config: {self.config}")
+        df = self.training_data.as_dataframe(flattened=True)
+        logger.info(f"Training data shape: {df.shape}")
+        target_cols = self.config.target_attributes
+        feature_cols = self.config.feature_attributes
+        if len(target_cols) != 1:
+            raise ValueError("Only one target column is supported")
+        if not feature_cols:
+            feature_cols = df.columns.difference(target_cols).tolist()
+            self.config.feature_attributes = feature_cols
+        target_col = target_cols[0]
+        logger.info(f"Feature columns: {feature_cols}")
+        X = df[feature_cols].copy()
+        logger.info(f"Target column: {target_col}")
+        y = df[target_col].copy()
+
+        # find list of features to skip (categorical with > N categories)
+        skip_features = []
+        for col in X.columns:
+            unique_values = self._get_unique_values(X[col])
+            if len(unique_values) > self.maximum_proportion_distinct_features * len(X[col]):
+                skip_features.append(col)
+            if False and (X[col].dtype == "object" or X[col].dtype.name == "category"):
+                if len(X[col].unique()) > self.maximum_proportion_distinct_features * len(X[col]):
+                    skip_features.append(col)
+        self.skip_features = skip_features
+        X = X.drop(skip_features, axis=1)
+        logger.info(f"Skipping features: {skip_features}")
+
+        # Encode features
+        encoded_features = []
+        for col in X.columns:
+            logger.info(f"Checking whether to encode: {col}")
+            col_encoder = self._get_encoder(X[col])
+            if col_encoder:
+                self.encoders[col] = col_encoder
+                if isinstance(col_encoder, OneHotEncoder):
+                    encoded = col_encoder.fit_transform(X[[col]])
+                    feature_names = col_encoder.get_feature_names_out([col])
+                    encoded_df = pd.DataFrame(encoded, columns=feature_names, index=X.index)
+                    X = pd.concat([X.drop(col, axis=1), encoded_df], axis=1)
+                    encoded_features.extend(feature_names)
+                elif isinstance(col_encoder, MultiLabelBinarizer):
+                    encoded = col_encoder.fit_transform(X[col])
+                    feature_names = [f"{col}_{c}" for c in col_encoder.classes_]
+                    encoded_df = pd.DataFrame(encoded, columns=feature_names, index=X.index)
+                    X = pd.concat([X.drop(col, axis=1), encoded_df], axis=1)
+                    encoded_features.extend(feature_names)
+                else:
+                    X[col] = col_encoder.fit_transform(X[col])
+                    encoded_features.append(col)
+            else:
+                encoded_features.append(col)
+
+        self.transformed_features = encoded_features
+        logger.info(f"Encoded features: {self.transformed_features}")
+        logger.info(f"Number of features after encoding: {len(self.transformed_features)}")
+
+        # Encode target
+        # y_encoder = LabelEncoder()
+        y_encoder = self._get_encoder(y)
+        if isinstance(y_encoder, OneHotEncoder):
+            y_encoder = LabelEncoder()
+        # self.encoders[target_col] = y_encoder
+        if y_encoder:
+            self.encoders[target_col] = y_encoder
+            y = y_encoder.fit_transform(y.values.ravel())  # Convert to 1D numpy array
+            self.transformed_targets = y_encoder.classes_
+
+        logger.info(f"Fitting model with features: {X.columns}")
+        clf = DecisionTreeClassifier(random_state=42)
+        clf.fit(X, y)
+        self.classifier = clf
+        logger.info("Model fit complete")
+        cv_scores = cross_val_score(self.classifier, X, y, cv=5)
+        self.confidence = cv_scores.mean()
+        logger.info(f"Cross-validation scores: {cv_scores}")
+
+    def derive(self, object: OBJECT) -> Optional[Inference]:
+        object = self._normalize(object)
+        new_X = pd.DataFrame([object])
+
+        # Apply encodings
+        encoded_features = {}
+        for col in self.config.feature_attributes:
+            if col in self.skip_features:
+                continue
+            if col in self.encoders:
+                encoder = self.encoders[col]
+                if isinstance(encoder, OneHotEncoder):
+                    print(f"Encoding: {col} v={object[col]} df={new_X[[col]]} encoder={encoder}")
+                    encoded = encoder.transform(new_X[[col]])
+                    feature_names = encoder.get_feature_names_out([col])
+                    for i, name in enumerate(feature_names):
+                        encoded_features[name] = encoded[0, i]
+                elif isinstance(encoder, MultiLabelBinarizer):
+                    encoded = encoder.transform(new_X[col])
+                    feature_names = [f"{col}_{c}" for c in encoder.classes_]
+                    for i, name in enumerate(feature_names):
+                        encoded_features[name] = encoded[0, i]
+                else:  # LabelEncoder or similar
+                    encoded_features[col] = encoder.transform(new_X[col].astype(str))[0]
+            else:
+                encoded_features[col] = new_X[col].iloc[0]
+
+        # Ensure all expected features are present and in the correct order
+        final_features = []
+        for feature in self.transformed_features:
+            if feature in encoded_features:
+                final_features.append(encoded_features[feature])
+            else:
+                final_features.append(0)  # or some other default value
+
+        # Create the final input array
+        new_X_array = np.array(final_features).reshape(1, -1)
+
+        logger.info(f"Input features: {self.transformed_features}")
+        logger.info(f"Number of input features: {len(self.transformed_features)}")
+
+        predictions = self.classifier.predict(new_X_array)
+        target_attribute = self.config.target_attributes[0]
+        y_encoder = self.encoders.get(target_attribute)
+
+        if y_encoder:
+            v = y_encoder.inverse_transform(predictions)
+        else:
+            v = predictions
+
+        predicted_object = {target_attribute: v[0]}
+        logger.info(f"Predicted object: {predicted_object}")
+        return Inference(predicted_object=predicted_object, confidence=self.confidence)
+
+    def _normalize(self, object: OBJECT) -> OBJECT:
+        """
+        Normalize the input object to ensure it has all the expected attributes.
+
+        Also remove any numpy/pandas oddities
+
+        :param object:
+        :return:
+        """
+        np_map = {np.nan: None}
+
+        def _tr(x: Any):
+            # TODO: figure a more elegant way to do this
+            try:
+                return np_map.get(x, x)
+            except TypeError:
+                return x
+
+        return {k: _tr(object.get(k, None)) for k in self.config.feature_attributes}
+
+    def export_model(
+        self, output: Optional[Union[str, Path, TextIO]], model_serialization: ModelSerialization = None, **kwargs
+    ):
+        def as_file():
+            if isinstance(output, (str, Path)):
+                return open(output, "w")
+            return output
+
+        if model_serialization is None:
+            if isinstance(output, (str, Path)):
+                model_serialization = ModelSerialization.from_filepath(output)
+            if model_serialization is None:
+                model_serialization = ModelSerialization.JOBLIB
+
+        if model_serialization == ModelSerialization.LINKML_EXPRESSION:
+            expr = tree_to_nested_expression(
+                self.classifier,
+                self.transformed_features,
+                self.encoders.keys(),
+                feature_encoders=self.encoders,
+                target_encoder=self.encoders.get(self.config.target_attributes[0]),
+            )
+            as_file().write(expr)
+        elif model_serialization == ModelSerialization.JOBLIB:
+            self.save_model(output)
+        elif model_serialization == ModelSerialization.RULE_BASED:
+            rbie = RuleBasedInferenceEngine(config=self.config)
+            rbie.import_model_from(self)
+            rbie.save_model(output)
+        elif model_serialization == ModelSerialization.PNG:
+            visualize_decision_tree(self.classifier, self.transformed_features, self.transformed_targets, output)
+        else:
+            raise ValueError(f"Unsupported model serialization: {model_serialization}")
+
+    def save_model(self, output: Union[str, Path]) -> None:
+        """
+        Save the trained model and related data to a file.
+
+        :param output: Path to save the model
+        """
+        import joblib
+
+        if self.classifier is None:
+            raise ValueError("Model has not been trained. Call initialize_model() first.")
+
+        # Use self.PERSIST_COLS
+        model_data = {k: getattr(self, k) for k in self.PERSIST_COLS}
+
+        joblib.dump(model_data, output)
+
+    @classmethod
+    def load_model(cls, file_path: Union[str, Path]) -> "SklearnInferenceEngine":
+        """
+        Load a trained model and related data from a file.
+
+        :param file_path: Path to the saved model
+        :return: SklearnInferenceEngine instance with loaded model
+        """
+        import joblib
+
+        model_data = joblib.load(file_path)
+
+        engine = cls(config=model_data["config"])
+        for k, v in model_data.items():
+            if k == "config":
+                continue
+            setattr(engine, k, v)
+
+        logger.info(f"Model loaded from {file_path}")
+        return engine

linkml_store/inference/inference_config.py

@@ -0,0 +1,62 @@
+import logging
+from typing import List, Optional, Tuple
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from linkml_store.api.collection import OBJECT
+from linkml_store.utils.format_utils import Format, load_objects
+
+logger = logging.getLogger(__name__)
+
+
+class LLMConfig(BaseModel, extra="forbid"):
+    """
+    Configuration for the LLM indexer.
+    """
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    model_name: str = "gpt-4o-mini"
+    token_limit: Optional[int] = None
+    number_of_few_shot_examples: Optional[int] = None
+    role: str = "Domain Expert"
+    cached_embeddings_database: Optional[str] = None
+    cached_embeddings_collection: Optional[str] = None
+    text_template: Optional[str] = None
+    text_template_syntax: Optional[str] = None
+
+
+class InferenceConfig(BaseModel, extra="forbid"):
+    """
+    Configuration for inference engines.
+    """
+
+    target_attributes: Optional[List[str]] = None
+    feature_attributes: Optional[List[str]] = None
+    train_test_split: Optional[Tuple[float, float]] = None
+    llm_config: Optional[LLMConfig] = None
+
+    @classmethod
+    def from_file(cls, file_path: str, format: Optional[Format] = None) -> "InferenceConfig":
+        """
+        Load an inference config from a file.
+
+        :param file_path: Path to the file.
+        :param format: Format of the file (YAML is recommended).
+        :return: InferenceConfig
+        """
+        if format and format.is_xsv():
+            logger.warning("XSV format is not recommended for inference config files")
+        objs = load_objects(file_path, format=format)
+        if len(objs) != 1:
+            raise ValueError(f"Expected 1 object, got {len(objs)}")
+        return cls(**objs[0])
+
+
+class Inference(BaseModel, extra="forbid"):
+    """
+    Result of an inference derivation.
+    """
+
+    predicted_object: OBJECT = Field(..., description="The predicted object.")
+    confidence: Optional[float] = Field(default=None, description="The confidence of the prediction.", le=1.0, ge=0.0)

linkml_store/inference/inference_engine.py

@@ -0,0 +1,200 @@
+import logging
+import random
+from abc import ABC
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import 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"
+
+    @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)
+
+    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
+
+    @property
+    def collection(self) -> Collection:
+        if not self._collection:
+            rows = self.base_collection.find({}, limit=-1).rows
+            # subset based on indices
+            subset = [rows[i] for i in self.indices]
+            db = self.base_collection.parent
+            subset_name = f"{self.base_collection.alias}__rag_{self.name}"
+            subset_collection = db.get_collection(subset_name, create_if_not_exists=True)
+            subset_collection.insert(subset)
+            self._collection = subset_collection
+        return self._collection
+
+    def as_dataframe(self, flattened=False) -> pd.DataFrame:
+        """
+        Return the slice of the collection as a dataframe.
+
+        :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:
+        :return:
+        """
+        split = split or self.config.train_test_split
+        if not split:
+            split = (0.7, 0.3)
+        logger.info(f"Loading and splitting data from collection {collection.alias}")
+        size = collection.size()
+        indices = range(size)
+        if randomize:
+            train_indices = 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)
+        # all_data = collection.find({}, limit=size).rows
+        # all_data_df = nested_objects_to_dataframe(all_data)
+        # all_data_df = collection.find({}, limit=size).rows_dataframe
+        # randomize/shuffle order of rows in dataframe
+        # all_data_df = all_data_df.sample(frac=1).reset_index(drop=True)
+        # self.training_data = CollectionSlice(dataframe=all_data_df[: int(size * split[0])])
+        # self.testing_data = CollectionSlice(dataframe=all_data_df[int(size * split[0]) : size])
+        # self.training_data = CollectionSlice(base_collection=collection, slice=(0, int(size * split[0])))
+        # self.testing_data = CollectionSlice(base_collection=collection, slice=(int(size * split[0]), size))
+
+    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)