orca-sdk 0.1.11__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/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/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/async_client.py

@@ -85,7 +85,7 @@ class BaseLabelPredictionResult(TypedDict):
     anomaly_score: float | None
     label: int | None
     label_name: str | None
-    logits: list[float]
+    logits: list[float] | None
 
 
 class BaseModel(TypedDict):
@@ -160,6 +160,18 @@ The type of a column in a datasource
 """
 
 
+class ComputeClassificationMetricsRequest(TypedDict):
+    expected_labels: list[int]
+    logits: list[list[float] | None]
+    anomaly_scores: NotRequired[list[float] | None]
+
+
+class ComputeRegressionMetricsRequest(TypedDict):
+    expected_scores: list[float]
+    predicted_scores: list[float | None]
+    anomaly_scores: NotRequired[list[float] | None]
+
+
 class ConstraintViolationErrorResponse(TypedDict):
     status_code: Literal[409]
     constraint: str
@@ -322,6 +334,7 @@ class GetDatasourceRowsRequest(TypedDict):
 
 class GetMemoriesRequest(TypedDict):
     memory_ids: list[str]
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
 
 
 class HealthyResponse(TypedDict):
@@ -392,6 +405,7 @@ class ListMemoriesRequest(TypedDict):
     offset: NotRequired[int]
     limit: NotRequired[int]
     filters: NotRequired[list[FilterItem]]
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
 
 
 class LookupRequest(TypedDict):
@@ -400,6 +414,7 @@ class LookupRequest(TypedDict):
     prompt: NotRequired[str | None]
     partition_id: NotRequired[str | list[str | None] | None]
     partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
 
 
 class LookupScoreMetrics(TypedDict):
@@ -570,8 +585,17 @@ class OrgPlan(TypedDict):
 
 class PRCurve(TypedDict):
     thresholds: list[float]
+    """
+    Threshold values for the curve
+    """
     precisions: list[float]
+    """
+    Precision values at each threshold
+    """
     recalls: list[float]
+    """
+    Recall values at each threshold
+    """
 
 
 class PredictionFeedback(TypedDict):
@@ -642,8 +666,17 @@ RARHeadType: TypeAlias = Literal["MMOE", "KNN"]
 
 class ROCCurve(TypedDict):
     thresholds: list[float]
+    """
+    Threshold values for the curve
+    """
     false_positive_rates: list[float]
+    """
+    False positive rate values at each threshold
+    """
     true_positive_rates: list[float]
+    """
+    True positive rate values at each threshold
+    """
 
 
 class ReadyResponse(TypedDict):
@@ -666,15 +699,49 @@ class RegressionEvaluationRequest(TypedDict):
 
 class RegressionMetrics(TypedDict):
     coverage: float
+    """
+    Percentage of predictions that are not none
+    """
     mse: float
+    """
+    Mean squared error of the predictions
+    """
     rmse: float
+    """
+    Root mean squared error of the predictions
+    """
     mae: float
+    """
+    Mean absolute error of the predictions
+    """
     r2: float
+    """
+    R-squared score (coefficient of determination) of the predictions
+    """
     explained_variance: float
+    """
+    Explained variance score of the predictions
+    """
     loss: float
+    """
+    Mean squared error loss of the predictions
+    """
     anomaly_score_mean: NotRequired[float | None]
+    """
+    Mean of anomaly scores across the dataset
+    """
     anomaly_score_median: NotRequired[float | None]
+    """
+    Median of anomaly scores across the dataset
+    """
     anomaly_score_variance: NotRequired[float | None]
+    """
+    Variance of anomaly scores across the dataset
+    """
+    warnings: NotRequired[list[str]]
+    """
+    Human-readable warnings about skipped or adjusted metrics
+    """
 
 
 class RegressionModelMetadata(TypedDict):
@@ -703,7 +770,7 @@ class RegressionPredictionRequest(TypedDict):
     save_telemetry_synchronously: NotRequired[bool]
     prompt: NotRequired[str | None]
     use_lookup_cache: NotRequired[bool]
-    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"] | None]
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
     ignore_unlabeled: NotRequired[bool]
     partition_ids: NotRequired[str | list[str | None] | None]
     partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
@@ -927,6 +994,7 @@ class GetMemorysetByNameOrIdMemoryByMemoryIdParams(TypedDict):
     """
     ID of the memory
     """
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
 
 
 class DeleteMemorysetByNameOrIdMemoryByMemoryIdParams(TypedDict):
@@ -1304,18 +1372,57 @@ class BootstrapLabeledMemoryDataResult(TypedDict):
 
 class ClassificationMetrics(TypedDict):
     coverage: float
+    """
+    Percentage of predictions that are not none
+    """
     f1_score: float
+    """
+    F1 score of the predictions
+    """
     accuracy: float
+    """
+    Accuracy of the predictions
+    """
     loss: float | None
+    """
+    Cross-entropy loss of the logits
+    """
     anomaly_score_mean: NotRequired[float | None]
+    """
+    Mean of anomaly scores across the dataset
+    """
     anomaly_score_median: NotRequired[float | None]
+    """
+    Median of anomaly scores across the dataset
+    """
     anomaly_score_variance: NotRequired[float | None]
+    """
+    Variance of anomaly scores across the dataset
+    """
     roc_auc: NotRequired[float | None]
+    """
+    Receiver operating characteristic area under the curve
+    """
     pr_auc: NotRequired[float | None]
+    """
+    Average precision (area under the curve of the precision-recall curve)
+    """
     pr_curve: NotRequired[PRCurve | None]
+    """
+    Precision-recall curve
+    """
     roc_curve: NotRequired[ROCCurve | None]
+    """
+    Receiver operating characteristic curve
+    """
     confusion_matrix: NotRequired[list[list[int]] | None]
+    """
+    Confusion matrix where the entry at row i, column j is the count of samples with true label i predicted as label j
+    """
     warnings: NotRequired[list[str]]
+    """
+    Human-readable warnings about skipped or adjusted metrics
+    """
 
 
 class ClassificationModelMetadata(TypedDict):
@@ -1348,7 +1455,7 @@ class ClassificationPredictionRequest(TypedDict):
     save_telemetry_synchronously: NotRequired[bool]
     prompt: NotRequired[str | None]
     use_lookup_cache: NotRequired[bool]
-    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"] | None]
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
     ignore_unlabeled: NotRequired[bool]
     partition_ids: NotRequired[str | list[str | None] | None]
     partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
@@ -1362,6 +1469,7 @@ class CloneMemorysetRequest(TypedDict):
     finetuned_embedding_model_name_or_id: NotRequired[str | None]
     max_seq_length_override: NotRequired[int | None]
     prompt: NotRequired[str]
+    is_partitioned: NotRequired[bool | None]
 
 
 class ColumnInfo(TypedDict):
@@ -1409,6 +1517,7 @@ class CreateMemorysetFromDatasourceRequest(TypedDict):
     prompt: NotRequired[str]
     hidden: NotRequired[bool]
     memory_type: NotRequired[MemoryType | None]
+    is_partitioned: NotRequired[bool]
     datasource_name_or_id: str
     datasource_label_column: NotRequired[str | None]
     datasource_score_column: NotRequired[str | None]
@@ -1433,6 +1542,7 @@ class CreateMemorysetRequest(TypedDict):
     prompt: NotRequired[str]
     hidden: NotRequired[bool]
     memory_type: NotRequired[MemoryType | None]
+    is_partitioned: NotRequired[bool]
 
 
 class CreateRegressionModelRequest(TypedDict):
@@ -1590,7 +1700,7 @@ class LabelPredictionWithMemoriesAndFeedback(TypedDict):
     anomaly_score: float | None
     label: int | None
     label_name: str | None
-    logits: list[float]
+    logits: list[float] | None
     timestamp: str
     input_value: str | bytes
     input_embedding: list[float]
@@ -1746,6 +1856,7 @@ class TelemetryMemoriesRequest(TypedDict):
     limit: NotRequired[int]
     filters: NotRequired[list[FilterItem | TelemetryFilterItem]]
     sort: NotRequired[list[TelemetrySortOptions] | None]
+    consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"]]
 
 
 class WorkerInfo(TypedDict):
@@ -1812,6 +1923,7 @@ class MemorysetMetadata(TypedDict):
     document_prompt_override: str | None
     query_prompt_override: str | None
     hidden: bool
+    is_partitioned: bool
     insertion_task_id: str | None
 
 
@@ -3660,6 +3772,46 @@ class OrcaAsyncClient(AsyncClient):
     ) -> EvaluationResponse:
         pass
 
+    @overload
+    async def POST(
+        self,
+        path: Literal["/classification_model/metrics"],
+        *,
+        params: None = None,
+        json: ComputeClassificationMetricsRequest,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> ClassificationMetrics:
+        pass
+
+    @overload
+    async def POST(
+        self,
+        path: Literal["/regression_model/metrics"],
+        *,
+        params: None = None,
+        json: ComputeRegressionMetricsRequest,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> RegressionMetrics:
+        pass
+
     @overload
     async def POST(
         self,