orca-sdk 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

orca_sdk/__init__.py

@@ -2,7 +2,7 @@
 OrcaSDK is a Python library for building and using retrieval augmented models in the OrcaCloud.
 """
 
-from ._utils.common import UNSET, CreateMode, DropMode
+from ._utils.common import UNSET, CreateMode, DropMode, logger
 from .classification_model import ClassificationMetrics, ClassificationModel
 from .client import OrcaClient
 from .credentials import OrcaCredentials
@@ -23,8 +23,8 @@ from .memoryset import (
     ScoredMemoryLookup,
     ScoredMemoryset,
 )
-from .regression_model import RegressionModel
+from .regression_model import RegressionMetrics, RegressionModel
 from .telemetry import ClassificationPrediction, FeedbackCategory, RegressionPrediction
 
 # only specify things that should show up on the root page of the reference docs because they are in private modules
-__all__ = ["UNSET", "CreateMode", "DropMode"]
+__all__ = ["UNSET", "CreateMode", "DropMode", "logger"]

orca_sdk/_utils/analysis_ui.py

@@ -5,7 +5,10 @@ import re
 from pathlib import Path
 from typing import TypedDict, cast
 
-import gradio as gr
+try:
+    import gradio as gr  # type: ignore
+except ImportError as e:
+    raise ImportError("gradio is required for UI features. Install it with: pip install orca_sdk[ui]") from e
 
 from ..memoryset import LabeledMemory, LabeledMemoryset
 

orca_sdk/_utils/auth.py

@@ -1,13 +1,12 @@
 """This module contains internal utils for managing api keys in tests"""
 
-import logging
 import os
 from typing import List, Literal
 
 from dotenv import load_dotenv
 
 from ..client import ApiKeyMetadata, OrcaClient
-from .common import DropMode
+from .common import DropMode, logger
 
 load_dotenv()  # this needs to be here to ensure env is populated before accessing it
 
@@ -59,7 +58,7 @@ def _authenticate_local_api(org_id: str = _DEFAULT_ORG_ID, api_key_name: str = "
     client = OrcaClient._resolve_client()
     client.base_url = "http://localhost:1584"
     client.headers.update({"Api-Key": _create_api_key(org_id, api_key_name)})
-    logging.info(f"Authenticated against local API at 'http://localhost:1584' with '{api_key_name}' API key")
+    logger.info(f"Authenticated against local API at 'http://localhost:1584' with '{api_key_name}' API key")
 
 
 __all__ = ["_create_api_key", "_delete_api_key", "_delete_org", "_list_api_keys", "_authenticate_local_api"]

orca_sdk/_utils/common.py

@@ -1,4 +1,21 @@
-from typing import Any, Literal
+import logging
+from typing import Any, Iterable, Iterator, Literal, TypeVar
+
+try:
+    from itertools import batched
+except ImportError:
+    # Polyfill for Python <3.12
+
+    from itertools import islice
+
+    _BatchT = TypeVar("_BatchT")
+
+    def batched(iterable: Iterable[_BatchT], n: int) -> Iterator[tuple[_BatchT, ...]]:
+        """Batch an iterable into chunks of size n (backfill for Python <3.12)."""
+        it = iter(iterable)
+        while batch := tuple(islice(it, n)):
+            yield batch
+
 
 CreateMode = Literal["error", "open"]
 """
@@ -35,3 +52,9 @@ UNSET: Any = _UnsetSentinel()
 """
 Default value to indicate that no update should be applied to a field and it should not be set to None
 """
+
+logger = logging.getLogger("orca_sdk")
+"""
+Logger for the Orca SDK.
+"""
+logger.addHandler(logging.NullHandler())

orca_sdk/_utils/prediction_result_ui.py

@@ -5,7 +5,10 @@ import re
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-import gradio as gr
+try:
+    import gradio as gr  # type: ignore
+except ImportError as e:
+    raise ImportError("gradio is required for UI features. Install it with: pip install orca_sdk[ui]") from e
 
 from ..memoryset import LabeledMemoryLookup, LabeledMemoryset, ScoredMemoryLookup
 

orca_sdk/_utils/torch_parsing.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from dataclasses import asdict, is_dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    # peer dependencies that are used for types only
+    from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore
+    from torch.utils.data import Dataset as TorchDataset  # type: ignore
+
+
+def parse_dict_like(item: Any, column_names: list[str] | None = None) -> dict:
+    if isinstance(item, dict):
+        return item
+
+    if isinstance(item, tuple):
+        if column_names is not None:
+            if len(item) != len(column_names):
+                raise ValueError(
+                    f"Tuple length ({len(item)}) does not match number of column names ({len(column_names)})"
+                )
+            return {column_names[i]: item[i] for i in range(len(item))}
+        elif hasattr(item, "_fields") and all(isinstance(field, str) for field in item._fields):  # type: ignore
+            return {field: getattr(item, field) for field in item._fields}  # type: ignore
+        else:
+            raise ValueError("For datasets that return unnamed tuples, please provide column_names argument")
+
+    if is_dataclass(item) and not isinstance(item, type):
+        return asdict(item)
+
+    raise ValueError(f"Cannot parse {type(item)}")
+
+
+def parse_batch(batch: Any, column_names: list[str] | None = None) -> list[dict]:
+    if isinstance(batch, list):
+        return [parse_dict_like(item, column_names) for item in batch]
+
+    batch = parse_dict_like(batch, column_names)
+    keys = list(batch.keys())
+    batch_size = len(batch[keys[0]])
+    for key in keys:
+        if not len(batch[key]) == batch_size:
+            raise ValueError(f"Batch must consist of values of the same length, but {key} has length {len(batch[key])}")
+    return [{key: batch[key][idx] for key in keys} for idx in range(batch_size)]
+
+
+def list_from_torch(
+    torch_data: TorchDataLoader | TorchDataset,
+    column_names: list[str] | None = None,
+) -> list[dict]:
+    """
+    Convert a PyTorch DataLoader or Dataset to a list of dictionaries.
+
+    Params:
+        torch_data: A PyTorch DataLoader or Dataset object to convert.
+        column_names: Optional list of column names to use for the data. If not provided,
+            the column names will be inferred from the data.
+    Returns:
+        A list of dictionaries containing the data from the PyTorch DataLoader or Dataset.
+    """
+    # peer dependency that is guaranteed to exist if the user provided a torch dataset
+    from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore
+
+    if isinstance(torch_data, TorchDataLoader):
+        dataloader = torch_data
+    else:
+        dataloader = TorchDataLoader(torch_data, batch_size=1, collate_fn=lambda x: x)
+
+    # Collect data from the dataloader into a list
+    data_list = []
+    try:
+        for batch in dataloader:
+            data_list.extend(parse_batch(batch, column_names=column_names))
+    except ValueError as e:
+        raise ValueError(str(e)) from e
+
+    return data_list

orca_sdk/_utils/torch_parsing_test.py

@@ -0,0 +1,142 @@
+from collections import namedtuple
+from dataclasses import dataclass
+
+import pytest
+
+from .torch_parsing import list_from_torch
+
+pytest.importorskip("torch")
+
+from torch.utils.data import DataLoader as TorchDataLoader  # noqa: E402
+from torch.utils.data import Dataset as TorchDataset  # noqa: E402
+
+
+def test_list_from_torch_dict_dataset(data: list[dict]):
+    class PytorchDictDataset(TorchDataset):
+        def __init__(self):
+            self.data = data
+
+        def __getitem__(self, i):
+            return self.data[i]
+
+        def __len__(self):
+            return len(self.data)
+
+    dataset = PytorchDictDataset()
+    data_list = list_from_torch(dataset)
+
+    assert isinstance(data_list, list)
+    assert len(data_list) == len(dataset)
+    assert set(list(data_list[0].keys())) == {"value", "label", "key", "score", "source_id", "partition_id"}
+
+
+def test_list_from_torch_dataloader(data: list[dict]):
+    class PytorchDictDataset(TorchDataset):
+        def __init__(self):
+            self.data = data
+
+        def __getitem__(self, i):
+            return self.data[i]
+
+        def __len__(self):
+            return len(self.data)
+
+    dataset = PytorchDictDataset()
+
+    def collate_fn(x: list[dict]):
+        return {"value": [item["value"] for item in x], "label": [item["label"] for item in x]}
+
+    dataloader = TorchDataLoader(dataset, batch_size=3, collate_fn=collate_fn)
+    data_list = list_from_torch(dataloader)
+
+    assert isinstance(data_list, list)
+    assert len(data_list) == len(dataset)
+    assert list(data_list[0].keys()) == ["value", "label"]
+
+
+def test_list_from_torch_tuple_dataset(data: list[dict]):
+    class PytorchTupleDataset(TorchDataset):
+        def __init__(self):
+            self.data = data
+
+        def __getitem__(self, i):
+            return self.data[i]["value"], self.data[i]["label"]
+
+        def __len__(self):
+            return len(self.data)
+
+    dataset = PytorchTupleDataset()
+
+    # raises error if no column names are passed in
+    with pytest.raises(ValueError):
+        list_from_torch(dataset)
+
+    # raises error if not enough column names are passed in
+    with pytest.raises(ValueError):
+        list_from_torch(dataset, column_names=["value"])
+
+    # creates list if correct number of column names are passed in
+    data_list = list_from_torch(dataset, column_names=["value", "label"])
+    assert isinstance(data_list, list)
+    assert len(data_list) == len(dataset)
+    assert list(data_list[0].keys()) == ["value", "label"]
+
+
+def test_list_from_torch_named_tuple_dataset(data: list[dict]):
+    # Given a Pytorch dataset that returns a namedtuple for each item
+    DatasetTuple = namedtuple("DatasetTuple", ["value", "label"])
+
+    class PytorchNamedTupleDataset(TorchDataset):
+        def __init__(self):
+            self.data = data
+
+        def __getitem__(self, i):
+            return DatasetTuple(self.data[i]["value"], self.data[i]["label"])
+
+        def __len__(self):
+            return len(self.data)
+
+    dataset = PytorchNamedTupleDataset()
+    data_list = list_from_torch(dataset)
+    assert isinstance(data_list, list)
+    assert len(data_list) == len(dataset)
+    assert list(data_list[0].keys()) == ["value", "label"]
+
+
+def test_list_from_torch_dataclass_dataset(data: list[dict]):
+    @dataclass
+    class DatasetItem:
+        text: str
+        label: int
+
+    class PytorchDataclassDataset(TorchDataset):
+        def __init__(self):
+            self.data = data
+
+        def __getitem__(self, i):
+            return DatasetItem(text=self.data[i]["value"], label=self.data[i]["label"])
+
+        def __len__(self):
+            return len(self.data)
+
+    dataset = PytorchDataclassDataset()
+    data_list = list_from_torch(dataset)
+    assert isinstance(data_list, list)
+    assert len(data_list) == len(dataset)
+    assert list(data_list[0].keys()) == ["text", "label"]
+
+
+def test_list_from_torch_invalid_dataset(data: list[dict]):
+    class PytorchInvalidDataset(TorchDataset):
+        def __init__(self):
+            self.data = data
+
+        def __getitem__(self, i):
+            return [self.data[i]["value"], self.data[i]["label"]]
+
+        def __len__(self):
+            return len(self.data)
+
+    dataset = PytorchInvalidDataset()
+    with pytest.raises(ValueError):
+        list_from_torch(dataset)

orca_sdk/_utils/value_parser.py

@@ -1,27 +1,43 @@
+from __future__ import annotations
+
 import base64
 import io
-from typing import cast
+from typing import TYPE_CHECKING, Any
 
-import numpy as np
-from numpy.typing import NDArray
-from PIL import Image as pil
+if TYPE_CHECKING:
+    # peer dependencies that are used for types only
+    import numpy as np  # type: ignore
+    from numpy.typing import NDArray  # type: ignore
+    from PIL import Image as pil  # type: ignore
 
-ValueType = str | pil.Image | NDArray[np.float32]
-"""
-The type of a value in a memoryset
+    ValueType = str | pil.Image | NDArray[np.float32]
+    """
+    The type of a value in a memoryset
 
-- `str`: string
-- `pil.Image`: image
-- `NDArray[np.float32]`: univariate or multivariate timeseries
-"""
+    - `str`: string
+    - `pil.Image`: image
+    - `NDArray[np.float32]`: univariate or multivariate timeseries
+    """
+else:
+    ValueType = Any
 
 
 def decode_value(value: str) -> ValueType:
     if value.startswith("data:image"):
+        try:
+            from PIL import Image as pil  # type: ignore
+        except ImportError as e:
+            raise ImportError("Install Pillow to use image values") from e
+
         header, data = value.split(",", 1)
         return pil.open(io.BytesIO(base64.b64decode(data)))
 
     if value.startswith("data:numpy"):
+        try:
+            import numpy as np  # type: ignore
+        except ImportError as e:
+            raise ImportError("Install numpy to use timeseries values") from e
+
         header, data = value.split(",", 1)
         return np.load(io.BytesIO(base64.b64decode(data)))
 
@@ -29,17 +45,28 @@ def decode_value(value: str) -> ValueType:
 
 
 def encode_value(value: ValueType) -> str:
-    if isinstance(value, pil.Image):
-        header = f"data:image/{value.format.lower()};base64," if value.format else "data:image;base64,"
+    try:
+        from PIL import Image as pil  # type: ignore
+    except ImportError:
+        pil = None  # type: ignore[assignment]
+
+    try:
+        import numpy as np  # type: ignore
+    except ImportError:
+        np = None  # type: ignore[assignment]
+
+    if pil is not None and isinstance(value, pil.Image):
+        header = f"data:image/{value.format.lower()};base64," if value.format else "data:image;base64,"  # type: ignore[union-attr]
         buffer = io.BytesIO()
-        value.save(buffer, format=value.format)
+        value.save(buffer, format=value.format)  # type: ignore[union-attr]
         bytes = buffer.getvalue()
         return header + base64.b64encode(bytes).decode("utf-8")
 
-    if isinstance(value, np.ndarray):
-        header = f"data:numpy/{value.dtype.name};base64,"
+    if np is not None and isinstance(value, np.ndarray):
+        header = f"data:numpy/{value.dtype.name};base64,"  # type: ignore[union-attr]
         buffer = io.BytesIO()
         np.save(buffer, value)
         return header + base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-    return value
+    # Value is already a string, or an unhandled type (fall back to str conversion)
+    return value if isinstance(value, str) else str(value)

orca_sdk/_utils/value_parser_test.py

@@ -1,5 +1,4 @@
-import numpy as np
-from PIL import Image as pil
+import pytest
 
 from .value_parser import decode_value, encode_value
 
@@ -13,6 +12,7 @@ def test_string_parsing():
 
 
 def test_image_parsing():
+    pil = pytest.importorskip("PIL.Image")
     img = pil.new("RGB", (10, 10), color="red")
     img.format = "PNG"
 
@@ -22,10 +22,11 @@ def test_image_parsing():
 
     decoded = decode_value(encoded)
     assert isinstance(decoded, pil.Image)
-    assert decoded.size == img.size
+    assert decoded.size == img.size  # type: ignore[union-attr]
 
 
 def test_timeseries_parsing():
+    np = pytest.importorskip("numpy")
     timeseries = np.random.rand(20, 3).astype(np.float32)
 
     encoded = encode_value(timeseries)
@@ -34,6 +35,6 @@ def test_timeseries_parsing():
 
     decoded = decode_value(encoded)
     assert isinstance(decoded, np.ndarray)
-    assert decoded.shape == timeseries.shape
-    assert decoded.dtype == timeseries.dtype
+    assert decoded.shape == timeseries.shape  # type: ignore[union-attr]
+    assert decoded.dtype == timeseries.dtype  # type: ignore[union-attr]
     assert np.allclose(decoded, timeseries)