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

linkml_store/inference/inference_engine.py

@@ -0,0 +1,173 @@
+import logging
+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, Field
+
+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)
+
+    collection: Optional[Collection] = None
+    dataframe: Optional[pd.DataFrame] = None
+    slice: Tuple[Optional[int], Optional[int]] = Field(default=(None, None))
+
+    def as_dataframe(self, flattened=False) -> pd.DataFrame:
+        """
+        Return the slice of the collection as a dataframe.
+
+        :return:
+        """
+        if self.dataframe is not None:
+            df = self.dataframe
+            return df.iloc[self.slice[0] : self.slice[1]]
+        elif self.collection is not None:
+            rs = self.collection.find({}, offset=self.slice[0], limit=self.slice[1] - self.slice[0])
+            if flattened:
+                return nested_objects_to_dataframe(rs.rows)
+            else:
+                return rs.rows_dataframe
+        else:
+            raise ValueError("No dataframe or collection provided")
+
+
+@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):
+        """
+        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()
+        self.training_data = CollectionSlice(collection=collection, slice=(0, int(size * split[0])))
+        self.testing_data = CollectionSlice(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)

linkml_store/utils/format_utils.py

@@ -27,6 +27,7 @@ class Format(Enum):
     JSON = "json"
     JSONL = "jsonl"
     YAML = "yaml"
+    YAMLL = "yamll"
     TSV = "tsv"
     CSV = "csv"
     PYTHON = "python"
@@ -63,6 +64,9 @@ class Format(Enum):
     def is_dump_format(self):
         return self in [Format.SQLDUMP_DUCKDB, Format.SQLDUMP_POSTGRES, Format.DUMP_MONGODB]
 
+    def is_xsv(self):
+        return self in [Format.TSV, Format.CSV]
+
 
 def load_objects_from_url(
     url: str,
@@ -135,11 +139,14 @@ def load_objects(
     compression: Optional[str] = None,
     expected_type: Optional[Type] = None,
     header_comment_token: Optional[str] = None,
+    select_query: Optional[str] = None,
 ) -> List[Dict[str, Any]]:
     """
     Load objects from a file or archive in supported formats.
     For tgz archives, it processes all files and concatenates the results.
 
+    TODO: Add schema hints for CSV/TSV parsing.
+
     :param file_path: The path to the file or archive.
     :param format: The format of the file. Can be a Format enum or a string value.
     :param compression: The compression type. Supports 'gz' for gzip and 'tgz' for tar.gz.
@@ -177,98 +184,22 @@ def load_objects(
             all_objects = process_file(f, format, expected_type, header_comment_token)
 
     logger.debug(f"Loaded {len(all_objects)} objects from {file_path}")
+    if select_query:
+        import jsonpath_ng as jp
+
+        path_expr = jp.parse(select_query)
+        new_objs = []
+        for obj in all_objects:
+            for match in path_expr.find(obj):
+                logging.debug(f"Match: {match.value}")
+                if isinstance(match.value, list):
+                    new_objs.extend(match.value)
+                else:
+                    new_objs.append(match.value)
+        all_objects = new_objs
     return all_objects
 
 
-def xxxload_objects(
-    file_path: Union[str, Path],
-    format: Union[Format, str] = None,
-    compression: Optional[str] = None,
-    expected_type: Type = None,
-    header_comment_token: Optional[str] = None,
-) -> List[Dict[str, Any]]:
-    """
-    Load objects from a file in JSON, JSONLines, YAML, CSV, or TSV format.
-
-    >>> load_objects("tests/input/test_data/data.csv")
-    [{'id': '1', 'name': 'John', 'age': '30'},
-     {'id': '2', 'name': 'Alice', 'age': '25'}, {'id': '3', 'name': 'Bob', 'age': '35'}]
-
-    :param file_path: The path to the file.
-    :param format: The format of the file. Can be a Format enum or a string value.
-    :param expected_type: The target type to load the objects into, e.g. list
-    :return: A list of dictionaries representing the loaded objects.
-    """
-    if isinstance(format, str):
-        format = Format(format)
-
-    if isinstance(file_path, Path):
-        file_path = str(file_path)
-
-    if not format and (file_path.endswith(".parquet") or file_path.endswith(".pq")):
-        format = Format.PARQUET
-    if not format and file_path.endswith(".tsv"):
-        format = Format.TSV
-    if not format and file_path.endswith(".csv"):
-        format = Format.CSV
-    if not format and file_path.endswith(".py"):
-        format = Format.PYTHON
-
-    mode = "r"
-    if format == Format.PARQUET:
-        mode = "rb"
-
-    if file_path == "-":
-        # set file_path to be a stream from stdin
-        f = sys.stdin
-    else:
-        f = open(file_path, mode)
-
-    if format == Format.JSON or (not format and file_path.endswith(".json")):
-        objs = json.load(f)
-    elif format == Format.JSONL or (not format and file_path.endswith(".jsonl")):
-        objs = [json.loads(line) for line in f]
-    elif format == Format.YAML or (not format and (file_path.endswith(".yaml") or file_path.endswith(".yml"))):
-        if expected_type and expected_type == list:  # noqa E721
-            objs = list(yaml.safe_load_all(f))
-        else:
-            objs = yaml.safe_load(f)
-    elif format == Format.TSV or format == Format.CSV:
-        # Skip initial comment lines if comment_char is set
-        if header_comment_token:
-            # Store the original position
-            original_pos = f.tell()
-
-            # Read and store lines until we find a non-comment line
-            lines = []
-            for line in f:
-                if not line.startswith(header_comment_token):
-                    break
-                lines.append(line)
-
-            # Go back to the original position
-            f.seek(original_pos)
-
-            # Skip the comment lines we found
-            for _ in lines:
-                f.readline()
-        if format == Format.TSV:
-            reader = csv.DictReader(f, delimiter="\t")
-        else:
-            reader = csv.DictReader(f)
-        objs = list(reader)
-    elif format == Format.PARQUET:
-        import pyarrow.parquet as pq
-
-        table = pq.read_table(f)
-        objs = table.to_pandas().to_dict(orient="records")
-    else:
-        raise ValueError(f"Unsupported file format: {file_path}")
-    if not isinstance(objs, list):
-        objs = [objs]
-    return objs
-
-
 def write_output(
     data: Union[List[Dict[str, Any]], Dict[str, Any], pd.DataFrame],
     format: Union[Format, str] = Format.YAML,
@@ -329,7 +260,7 @@ def render_output(
     if format == Format.FORMATTED:
         if not isinstance(data, pd.DataFrame):
             data = pd.DataFrame(data)
-        return str(data)
+        return data.to_string(max_rows=None)
 
     if isinstance(data, pd.DataFrame):
         data = data.to_dict(orient="records")

linkml_store/utils/llm_utils.py

@@ -0,0 +1,95 @@
+from typing import Callable, List, Optional
+
+from tiktoken import Encoding
+
+MODEL_TOKEN_MAPPING = {
+    "gpt-4o-mini": 128_000,
+    "gpt-4o": 128_000,
+    "gpt-4o-2024-05-13": 128_000,
+    "gpt-4": 8192,
+    "gpt-4-0314": 8192,
+    "gpt-4-0613": 8192,
+    "gpt-4-32k": 32768,
+    "gpt-4-32k-0314": 32768,
+    "gpt-4-32k-0613": 32768,
+    "gpt-3.5-turbo": 4096,
+    "gpt-3.5-turbo-0301": 4096,
+    "gpt-3.5-turbo-0613": 4096,
+    "gpt-3.5-turbo-16k": 16385,
+    "gpt-3.5-turbo-16k-0613": 16385,
+    "gpt-3.5-turbo-instruct": 4096,
+    "text-ada-001": 2049,
+    "ada": 2049,
+    "text-babbage-001": 2040,
+    "babbage": 2049,
+    "text-curie-001": 2049,
+    "curie": 2049,
+    "davinci": 2049,
+    "text-davinci-003": 4097,
+    "text-davinci-002": 4097,
+    "code-davinci-002": 8001,
+    "code-davinci-001": 8001,
+    "code-cushman-002": 2048,
+    "code-cushman-001": 2048,
+    "claude": 200_000,
+}
+
+
+def render_formatted_text(
+    render_func: Callable,
+    values: List[str],
+    encoding: Encoding,
+    token_limit: int,
+    additional_text: Optional[str] = None,
+) -> str:
+    """
+    Render a formatted text string with a given object, encoding, and token limit.
+
+    >>> from tiktoken import encoding_for_model
+    >>> encoding = encoding_for_model("gpt-4o-mini")
+    >>> names = ["Alice", "Bob", "DoctorHippopotamusMcHippopotamusFace"]
+    >>> f = lambda x: f"Hello, {' '.join(x)}!"
+    >>> render_formatted_text(f, names, encoding, 4096)
+    'Hello, Alice Bob DoctorHippopotamusMcHippopotamusFace!'
+    >>> render_formatted_text(f, names, encoding, 5)
+    'Hello, Alice Bob!'
+
+    :param render_func: Rendering function
+    :param values: Values to render
+    :param encoding: Encoding
+    :param token_limit: Token limit
+    :param additional_text: Additional text to consider
+    :return:
+    """
+    text = render_func(values)
+    if additional_text:
+        token_limit -= len(encoding.encode(additional_text))
+    text_length = len(encoding.encode(text))
+    if text_length <= token_limit:
+        return text
+    if not values:
+        raise ValueError(f"Cannot fit text into token limit: {text_length} > {token_limit}")
+    return render_formatted_text(render_func, values[0:-1], encoding=encoding, token_limit=token_limit)
+
+
+def get_token_limit(model_name: str) -> int:
+    """
+    Estimate the token limit for a model.
+
+    >>> get_token_limit("gpt-4o-mini")
+    128000
+
+    also works with nested names:
+
+    >>> get_token_limit("my/claude-opus")
+    200000
+
+
+    :param model_name: Model name
+    :return: Estimated token limit
+    """
+    # sort MODEL_TOKEN_MAPPING by key length to ensure that the longest model names are checked first
+    for model, token_limit in sorted(MODEL_TOKEN_MAPPING.items(), key=lambda x: len(x[0]), reverse=True):
+        if model in model_name:
+            return token_limit
+    return 4096

linkml_store/utils/object_utils.py

@@ -29,7 +29,7 @@ def object_path_update(
     """
     if isinstance(obj, BaseModel):
         typ = type(obj)
-        obj = obj.dict()
+        obj = obj.model_dump(exclude_none=True)
         obj = object_path_update(obj, path, value)
         return typ(**obj)
     obj = deepcopy(obj)
@@ -45,6 +45,8 @@ def object_path_update(
                 obj.append({})
             obj = obj[index]
         else:
+            if part in obj and obj[part] is None:
+                del obj[part]
             obj = obj.setdefault(part, {})
     last_part = parts[-1]
     if "[" in last_part:

linkml_store/utils/pandas_utils.py

@@ -1,7 +1,59 @@
-from typing import Dict, List, Tuple, Union
+import logging
+from typing import Any, Dict, List, Tuple, Union
 
 import pandas as pd
 
+logger = logging.getLogger(__name__)
+
+
+def flatten_dict(d: Dict[str, Any], parent_key: str = "", sep: str = ".") -> Dict[str, Any]:
+    """
+    Recursively flatten a nested dictionary.
+
+    Args:
+        d (Dict[str, Any]): The dictionary to flatten.
+        parent_key (str): The parent key for nested dictionaries.
+        sep (str): The separator to use between keys.
+
+    Returns:
+        Dict[str, Any]: A flattened dictionary.
+
+    >>> flatten_dict({'a': 1, 'b': {'c': 2, 'd': {'e': 3}}})
+    {'a': 1, 'b.c': 2, 'b.d.e': 3}
+    """
+    items = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+def nested_objects_to_dataframe(data: List[Dict[str, Any]]) -> pd.DataFrame:
+    """
+    Convert a list of nested objects to a flattened pandas DataFrame.
+
+    Args:
+        data (List[Dict[str, Any]]): A list of nested dictionaries.
+
+    Returns:
+        pd.DataFrame: A flattened DataFrame.
+
+    >>> data = [
+    ...     {"person": {"name": "Alice", "age": 30}, "job": {"title": "Engineer", "salary": 75000}},
+    ...     {"person": {"name": "Bob", "age": 35}, "job": {"title": "Manager", "salary": 85000}}
+    ... ]
+    >>> df = nested_objects_to_dataframe(data)
+    >>> df.columns.tolist()
+    ['person.name', 'person.age', 'job.title', 'job.salary']
+    >>> df['person.name'].tolist()
+    ['Alice', 'Bob']
+    """
+    flattened_data = [flatten_dict(item) for item in data]
+    return pd.DataFrame(flattened_data)
+
 
 def facet_summary_to_dataframe_unmelted(
     facet_summary: Dict[Union[str, Tuple[str, ...]], List[Tuple[Union[str, Tuple[str, ...]], int]]]
@@ -22,7 +74,8 @@ def facet_summary_to_dataframe_unmelted(
                     categories, value = cat_val_tuple[:-1], cat_val_tuple[-1]
                 row = {"Value": value}
                 for i, facet in enumerate(facet_type):
-                    row[facet] = categories[i]
+                    logger.debug(f"FT={facet_type} i={i} Facet: {facet}, categories: {categories}")
+                    row[facet] = categories[i] if len(categories) > i else None
                 rows.append(row)
 
     df = pd.DataFrame(rows)