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

orca_sdk/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 OrcaClient(Client):
     ) -> EvaluationResponse:
         pass
 
+    @overload
+    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
+    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
     def POST(
         self,

orca_sdk/conftest.py

@@ -24,15 +24,6 @@ os.environ["ORCA_API_URL"] = os.environ.get("ORCA_API_URL", "http://localhost:15
 os.environ["ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY"] = "true"
 
 
-def skip_in_prod(reason: str):
-    """Custom decorator to skip tests when running against production API"""
-    PROD_API_URLs = ["https://api.orcadb.ai", "https://api.staging.orcadb.ai"]
-    return pytest.mark.skipif(
-        os.environ["ORCA_API_URL"] in PROD_API_URLs,
-        reason=reason,
-    )
-
-
 def skip_in_ci(reason: str):
     """Custom decorator to skip tests when running in CI"""
     return pytest.mark.skipif(
@@ -201,6 +192,11 @@ SAMPLE_DATA = [
 ]
 
 
+@pytest.fixture(scope="session")
+def data() -> list[dict]:
+    return SAMPLE_DATA
+
+
 @pytest.fixture(scope="session")
 def hf_dataset(label_names: list[str]) -> Dataset:
     return Dataset.from_list(
@@ -232,6 +228,11 @@ EVAL_DATASET = [
 ]
 
 
+@pytest.fixture(scope="session")
+def eval_data() -> list[dict]:
+    return EVAL_DATASET
+
+
 @pytest.fixture(scope="session")
 def eval_datasource() -> Datasource:
     eval_datasource = Datasource.from_list("eval_datasource", EVAL_DATASET)

orca_sdk/datasource.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import logging
 import tempfile
 import zipfile
 from datetime import datetime
@@ -9,17 +8,18 @@ from os import PathLike
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Union, cast
 
-from datasets import Dataset, DatasetDict
 from httpx._types import FileTypes  # type: ignore
 from tqdm.auto import tqdm
 
-from ._utils.common import CreateMode, DropMode
-from ._utils.data_parsing import hf_dataset_from_torch
+from ._utils.common import CreateMode, DropMode, logger
+from ._utils.torch_parsing import list_from_torch
 from ._utils.tqdm_file_reader import TqdmFileReader
 from .client import DatasourceMetadata, OrcaClient
 
 if TYPE_CHECKING:
     # These are peer dependencies that are used for types only
+    from datasets import Dataset as HFDataset  # type: ignore
+    from datasets import DatasetDict as HFDatasetDict  # type: ignore
     from pandas import DataFrame as PandasDataFrame  # type: ignore
     from pyarrow import Table as PyArrowTable  # type: ignore
     from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore
@@ -146,7 +146,7 @@ class Datasource:
 
     @classmethod
     def from_hf_dataset(
-        cls, name: str, dataset: Dataset, if_exists: CreateMode = "error", description: str | None = None
+        cls, name: str, dataset: HFDataset, if_exists: CreateMode = "error", description: str | None = None
     ) -> Datasource:
         """
         Create a new datasource from a Hugging Face Dataset
@@ -183,7 +183,7 @@ class Datasource:
     def from_hf_dataset_dict(
         cls,
         name: str,
-        dataset_dict: DatasetDict,
+        dataset_dict: HFDatasetDict,
         if_exists: CreateMode = "error",
         description: dict[str, str | None] | str | None = None,
     ) -> dict[str, Datasource]:
@@ -239,8 +239,8 @@ class Datasource:
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
         """
-        hf_dataset = hf_dataset_from_torch(torch_data, column_names=column_names)
-        return cls.from_hf_dataset(name, hf_dataset, if_exists=if_exists, description=description)
+        data_list = list_from_torch(torch_data, column_names=column_names)
+        return cls.from_list(name, data_list, if_exists=if_exists, description=description)
 
     @classmethod
     def from_list(
@@ -326,14 +326,24 @@ class Datasource:
                 `"error"`. Other option is `"open"` to open the existing datasource.
             description: Optional description for the datasource
 
+        Notes:
+            Data type precision may be lost during upload unless the [`datasets`][datasets] library is installed.
+
         Returns:
             A handle to the new datasource in the OrcaCloud
 
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
+            ImportError: If the upload dependency group is not installed
         """
-        dataset = Dataset.from_pandas(dataframe)
-        return cls.from_hf_dataset(name, dataset, if_exists=if_exists, description=description)
+        try:
+            from datasets import Dataset  # type: ignore
+
+            return cls.from_hf_dataset(
+                name, Dataset.from_pandas(dataframe), if_exists=if_exists, description=description
+            )
+        except ImportError:
+            return cls.from_dict(name, dataframe.to_dict(orient="list"), if_exists=if_exists, description=description)
 
     @classmethod
     def from_arrow(
@@ -404,6 +414,7 @@ class Datasource:
 
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
+            ImportError: If the path is a directory and [`datasets`][datasets] is not installed
         """
         # Check if datasource already exists and handle accordingly
         existing = _handle_existing_datasource(name, if_exists)
@@ -414,6 +425,13 @@ class Datasource:
 
         # For dataset directories, use the upload endpoint with multiple files
         if file_path.is_dir():
+            try:
+                from datasets import Dataset  # type: ignore
+            except ImportError as e:
+                raise ImportError(
+                    "The path is a directory, we only support uploading directories that contain saved HuggingFace datasets but datasets is not installed."
+                ) from e
+
             return cls.from_hf_dataset(
                 name, Dataset.load_from_disk(file_path), if_exists=if_exists, description=description
             )
@@ -484,7 +502,7 @@ class Datasource:
         try:
             client = OrcaClient._resolve_client()
             client.DELETE("/datasource/{name_or_id}", params={"name_or_id": name_or_id})
-            logging.info(f"Deleted datasource {name_or_id}")
+            logger.info(f"Deleted datasource {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
                 raise
@@ -566,9 +584,9 @@ class Datasource:
             with zipfile.ZipFile(output_path, "r") as zip_ref:
                 zip_ref.extractall(extract_dir)
             output_path.unlink()  # Remove the zip file after extraction
-            logging.info(f"Downloaded {extract_dir}")
+            logger.info(f"Downloaded {extract_dir}")
         else:
-            logging.info(f"Downloaded {output_path}")
+            logger.info(f"Downloaded {output_path}")
 
     def to_list(self) -> list[dict]:
         """

orca_sdk/embedding_model.py

@@ -4,8 +4,7 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import TYPE_CHECKING, Literal, Sequence, cast, get_args, overload
 
-from ._shared.metrics import ClassificationMetrics, RegressionMetrics
-from ._utils.common import UNSET, CreateMode, DropMode
+from ._utils.common import CreateMode, DropMode
 from .client import (
     EmbeddingEvaluationRequest,
     EmbeddingFinetuningMethod,
@@ -20,7 +19,9 @@ from .datasource import Datasource
 from .job import Job, Status
 
 if TYPE_CHECKING:
+    from .classification_model import ClassificationMetrics
     from .memoryset import LabeledMemoryset, ScoredMemoryset
+    from .regression_model import RegressionMetrics
 
 
 class EmbeddingModelBase(ABC):
@@ -230,6 +231,9 @@ class EmbeddingModelBase(ABC):
             raise ValueError("Invalid embedding model")
 
         def get_result(job_id: str) -> ClassificationMetrics | RegressionMetrics:
+            from .classification_model import ClassificationMetrics
+            from .regression_model import RegressionMetrics
+
             client = OrcaClient._resolve_client()
             if isinstance(self, PretrainedEmbeddingModel):
                 res = client.GET(
@@ -244,34 +248,7 @@ class EmbeddingModelBase(ABC):
             else:
                 raise ValueError("Invalid embedding model")
             assert res is not None
-            return (
-                RegressionMetrics(
-                    coverage=res.get("coverage"),
-                    mse=res.get("mse"),
-                    rmse=res.get("rmse"),
-                    mae=res.get("mae"),
-                    r2=res.get("r2"),
-                    explained_variance=res.get("explained_variance"),
-                    loss=res.get("loss"),
-                    anomaly_score_mean=res.get("anomaly_score_mean"),
-                    anomaly_score_median=res.get("anomaly_score_median"),
-                    anomaly_score_variance=res.get("anomaly_score_variance"),
-                )
-                if "mse" in res
-                else ClassificationMetrics(
-                    coverage=res.get("coverage"),
-                    f1_score=res.get("f1_score"),
-                    accuracy=res.get("accuracy"),
-                    loss=res.get("loss"),
-                    anomaly_score_mean=res.get("anomaly_score_mean"),
-                    anomaly_score_median=res.get("anomaly_score_median"),
-                    anomaly_score_variance=res.get("anomaly_score_variance"),
-                    roc_auc=res.get("roc_auc"),
-                    pr_auc=res.get("pr_auc"),
-                    pr_curve=res.get("pr_curve"),
-                    roc_curve=res.get("roc_curve"),
-                )
-            )
+            return RegressionMetrics(res) if "mse" in res else ClassificationMetrics(res)
 
         job = Job(response["job_id"], lambda: get_result(response["job_id"]))
         return job if background else job.result()
@@ -404,7 +381,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         return isinstance(other, PretrainedEmbeddingModel) and self.name == other.name
 
     def __repr__(self) -> str:
-        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {self.num_params/1000000:.0f}M}})"
+        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {self.num_params / 1000000:.0f}M}})"
 
     @classmethod
     def all(cls) -> list[PretrainedEmbeddingModel]:

orca_sdk/embedding_model_test.py

@@ -4,9 +4,9 @@ from uuid import uuid4
 
 import pytest
 
+from .classification_model import ClassificationMetrics
 from .datasource import Datasource
 from .embedding_model import (
-    ClassificationMetrics,
     FinetunedEmbeddingModel,
     PretrainedEmbeddingModel,
     PretrainedEmbeddingModelName,