orca-sdk 0.0.91__py3-none-any.whl → 0.0.92__py3-none-any.whl

orca_sdk/_generated_api_client/models/base_label_prediction_result.py

@@ -10,7 +10,7 @@ The main change is:
 
 # flake8: noqa: C901
 
-from typing import Any, Type, TypeVar, Union, cast
+from typing import Any, List, Type, TypeVar, Union, cast
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
@@ -28,6 +28,7 @@ class BaseLabelPredictionResult:
         anomaly_score (Union[None, float]):
         label (int):
         label_name (Union[None, str]):
+        logits (List[float]):
     """
 
     prediction_id: Union[None, str]
@@ -35,6 +36,7 @@ class BaseLabelPredictionResult:
     anomaly_score: Union[None, float]
     label: int
     label_name: Union[None, str]
+    logits: List[float]
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
@@ -51,6 +53,8 @@ class BaseLabelPredictionResult:
         label_name: Union[None, str]
         label_name = self.label_name
 
+        logits = self.logits
+
         field_dict: dict[str, Any] = {}
         field_dict.update(self.additional_properties)
         field_dict.update(
@@ -60,6 +64,7 @@ class BaseLabelPredictionResult:
                 "anomaly_score": anomaly_score,
                 "label": label,
                 "label_name": label_name,
+                "logits": logits,
             }
         )
 
@@ -94,12 +99,15 @@ class BaseLabelPredictionResult:
 
         label_name = _parse_label_name(d.pop("label_name"))
 
+        logits = cast(List[float], d.pop("logits"))
+
         base_label_prediction_result = cls(
             prediction_id=prediction_id,
             confidence=confidence,
             anomaly_score=anomaly_score,
             label=label,
             label_name=label_name,
+            logits=logits,
         )
 
         base_label_prediction_result.additional_properties = d

orca_sdk/_generated_api_client/models/label_prediction_with_memories_and_feedback.py

@@ -34,10 +34,10 @@ class LabelPredictionWithMemoriesAndFeedback:
         anomaly_score (Union[None, float]):
         label (int):
         label_name (Union[None, str]):
+        logits (List[float]):
         timestamp (datetime.datetime):
         input_value (str):
         input_embedding (List[float]):
-        logits (List[float]):
         expected_label (Union[None, int]):
         expected_label_name (Union[None, str]):
         memories (List['LabelPredictionMemoryLookup']):
@@ -56,10 +56,10 @@ class LabelPredictionWithMemoriesAndFeedback:
     anomaly_score: Union[None, float]
     label: int
     label_name: Union[None, str]
+    logits: List[float]
     timestamp: datetime.datetime
     input_value: str
     input_embedding: List[float]
-    logits: List[float]
     expected_label: Union[None, int]
     expected_label_name: Union[None, str]
     memories: List["LabelPredictionMemoryLookup"]
@@ -86,6 +86,8 @@ class LabelPredictionWithMemoriesAndFeedback:
         label_name: Union[None, str]
         label_name = self.label_name
 
+        logits = self.logits
+
         timestamp = self.timestamp.isoformat()
 
         input_value: str
@@ -93,8 +95,6 @@ class LabelPredictionWithMemoriesAndFeedback:
 
         input_embedding = self.input_embedding
 
-        logits = self.logits
-
         expected_label: Union[None, int]
         expected_label = self.expected_label
 
@@ -136,10 +136,10 @@ class LabelPredictionWithMemoriesAndFeedback:
                 "anomaly_score": anomaly_score,
                 "label": label,
                 "label_name": label_name,
+                "logits": logits,
                 "timestamp": timestamp,
                 "input_value": input_value,
                 "input_embedding": input_embedding,
-                "logits": logits,
                 "expected_label": expected_label,
                 "expected_label_name": expected_label_name,
                 "memories": memories,
@@ -182,6 +182,8 @@ class LabelPredictionWithMemoriesAndFeedback:
 
         label_name = _parse_label_name(d.pop("label_name"))
 
+        logits = cast(List[float], d.pop("logits"))
+
         timestamp = isoparse(d.pop("timestamp"))
 
         def _parse_input_value(data: object) -> str:
@@ -191,8 +193,6 @@ class LabelPredictionWithMemoriesAndFeedback:
 
         input_embedding = cast(List[float], d.pop("input_embedding"))
 
-        logits = cast(List[float], d.pop("logits"))
-
         def _parse_expected_label(data: object) -> Union[None, int]:
             if data is None:
                 return data
@@ -251,10 +251,10 @@ class LabelPredictionWithMemoriesAndFeedback:
             anomaly_score=anomaly_score,
             label=label,
             label_name=label_name,
+            logits=logits,
             timestamp=timestamp,
             input_value=input_value,
             input_embedding=input_embedding,
-            logits=logits,
             expected_label=expected_label,
             expected_label_name=expected_label_name,
             memories=memories,

orca_sdk/_generated_api_client/models/labeled_memoryset_metadata.py

@@ -43,6 +43,7 @@ class LabeledMemorysetMetadata:
         label_names (List[str]):
         created_at (datetime.datetime):
         updated_at (datetime.datetime):
+        memories_updated_at (datetime.datetime):
         insertion_task_id (str):
         insertion_status (TaskStatus): Status of task in the task queue
         metrics (MemorysetMetrics):
@@ -59,6 +60,7 @@ class LabeledMemorysetMetadata:
     label_names: List[str]
     created_at: datetime.datetime
     updated_at: datetime.datetime
+    memories_updated_at: datetime.datetime
     insertion_task_id: str
     insertion_status: TaskStatus
     metrics: "MemorysetMetrics"
@@ -97,6 +99,8 @@ class LabeledMemorysetMetadata:
 
         updated_at = self.updated_at.isoformat()
 
+        memories_updated_at = self.memories_updated_at.isoformat()
+
         insertion_task_id = self.insertion_task_id
 
         insertion_status = (
@@ -120,6 +124,7 @@ class LabeledMemorysetMetadata:
                 "label_names": label_names,
                 "created_at": created_at,
                 "updated_at": updated_at,
+                "memories_updated_at": memories_updated_at,
                 "insertion_task_id": insertion_task_id,
                 "insertion_status": insertion_status,
                 "metrics": metrics,
@@ -180,6 +185,8 @@ class LabeledMemorysetMetadata:
 
         updated_at = isoparse(d.pop("updated_at"))
 
+        memories_updated_at = isoparse(d.pop("memories_updated_at"))
+
         insertion_task_id = d.pop("insertion_task_id")
 
         insertion_status = TaskStatus(d.pop("insertion_status"))
@@ -198,6 +205,7 @@ class LabeledMemorysetMetadata:
             label_names=label_names,
             created_at=created_at,
             updated_at=updated_at,
+            memories_updated_at=memories_updated_at,
             insertion_task_id=insertion_task_id,
             insertion_status=insertion_status,
             metrics=metrics,

orca_sdk/_shared/__init__.py

@@ -0,0 +1 @@
+from .metrics import calculate_pr_curve, calculate_roc_curve, compute_classifier_metrics

orca_sdk/_shared/metrics.py

@@ -0,0 +1,195 @@
+"""
+This module contains metrics for usage with the Hugging Face Trainer.
+
+IMPORTANT:
+- This is a shared file between OrcaLib and the Orca SDK.
+- Please ensure that it does not have any dependencies on the OrcaLib code.
+- Make sure to edit this file in orcalib/shared and NOT in orca_sdk, since it will be overwritten there.
+
+"""
+
+from typing import Literal, Tuple, TypedDict
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.special import softmax
+from sklearn.metrics import accuracy_score, auc, f1_score, log_loss
+from sklearn.metrics import precision_recall_curve as sklearn_precision_recall_curve
+from sklearn.metrics import roc_auc_score
+from sklearn.metrics import roc_curve as sklearn_roc_curve
+from transformers.trainer_utils import EvalPrediction
+
+
+class ClassificationMetrics(TypedDict):
+    accuracy: float
+    f1_score: float
+    roc_auc: float | None  # receiver operating characteristic area under the curve (if all classes are present)
+    pr_auc: float | None  # precision-recall area under the curve (only for binary classification)
+    log_loss: float  # cross-entropy loss for probabilities
+
+
+def compute_classifier_metrics(eval_pred: EvalPrediction) -> ClassificationMetrics:
+    """
+    Compute standard metrics for classifier with Hugging Face Trainer.
+
+    Args:
+        eval_pred: The predictions containing logits and expected labels as given by the Trainer.
+
+    Returns:
+        A dictionary containing the accuracy, f1 score, and ROC AUC score.
+    """
+    logits, references = eval_pred
+    if isinstance(logits, tuple):
+        logits = logits[0]
+    if not isinstance(logits, np.ndarray):
+        raise ValueError("Logits must be a numpy array")
+    if not isinstance(references, np.ndarray):
+        raise ValueError(
+            "Multiple label columns found, use the `label_names` training argument to specify which one to use"
+        )
+
+    if not (logits > 0).all():
+        # convert logits to probabilities with softmax if necessary
+        probabilities = softmax(logits)
+    elif not np.allclose(logits.sum(-1, keepdims=True), 1.0):
+        # convert logits to probabilities through normalization if necessary
+        probabilities = logits / logits.sum(-1, keepdims=True)
+    else:
+        probabilities = logits
+
+    return classification_scores(references, probabilities)
+
+
+def classification_scores(
+    references: NDArray[np.int64],
+    probabilities: NDArray[np.float32],
+    average: Literal["micro", "macro", "weighted", "binary"] | None = None,
+    multi_class: Literal["ovr", "ovo"] = "ovr",
+) -> ClassificationMetrics:
+    if probabilities.ndim == 1:
+        # convert 1D probabilities (binary) to 2D logits
+        probabilities = np.column_stack([1 - probabilities, probabilities])
+    elif probabilities.ndim == 2:
+        if probabilities.shape[1] < 2:
+            raise ValueError("Use a different metric function for regression tasks")
+    else:
+        raise ValueError("Probabilities must be 1 or 2 dimensional")
+
+    predictions = np.argmax(probabilities, axis=-1)
+
+    num_classes_references = len(set(references))
+    num_classes_predictions = len(set(predictions))
+
+    if average is None:
+        average = "binary" if num_classes_references == 2 else "weighted"
+
+    accuracy = accuracy_score(references, predictions)
+    f1 = f1_score(references, predictions, average=average)
+    loss = log_loss(references, probabilities)
+
+    if num_classes_references == num_classes_predictions:
+        # special case for binary classification: https://github.com/scikit-learn/scikit-learn/issues/20186
+        if num_classes_references == 2:
+            roc_auc = roc_auc_score(references, probabilities[:, 1])
+            precisions, recalls, _ = calculate_pr_curve(references, probabilities[:, 1])
+            pr_auc = auc(recalls, precisions)
+        else:
+            roc_auc = roc_auc_score(references, probabilities, multi_class=multi_class)
+            pr_auc = None
+    else:
+        roc_auc = None
+        pr_auc = None
+
+    return {
+        "accuracy": float(accuracy),
+        "f1_score": float(f1),
+        "roc_auc": float(roc_auc) if roc_auc is not None else None,
+        "pr_auc": float(pr_auc) if pr_auc is not None else None,
+        "log_loss": float(loss),
+    }
+
+
+def calculate_pr_curve(
+    references: NDArray[np.int64],
+    probabilities: NDArray[np.float32],
+    max_length: int = 100,
+) -> Tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:
+    if probabilities.ndim == 1:
+        probabilities_slice = probabilities
+    elif probabilities.ndim == 2:
+        probabilities_slice = probabilities[:, 1]
+    else:
+        raise ValueError("Probabilities must be 1 or 2 dimensional")
+
+    if len(probabilities_slice) != len(references):
+        raise ValueError("Probabilities and references must have the same length")
+
+    precisions, recalls, thresholds = sklearn_precision_recall_curve(references, probabilities_slice)
+
+    # Convert all arrays to float32 immediately after getting them
+    precisions = precisions.astype(np.float32)
+    recalls = recalls.astype(np.float32)
+    thresholds = thresholds.astype(np.float32)
+
+    # Concatenate with 0 to include the lowest threshold
+    thresholds = np.concatenate(([0], thresholds))
+
+    # Sort by threshold
+    sorted_indices = np.argsort(thresholds)
+    thresholds = thresholds[sorted_indices]
+    precisions = precisions[sorted_indices]
+    recalls = recalls[sorted_indices]
+
+    if len(precisions) > max_length:
+        new_thresholds = np.linspace(0, 1, max_length, dtype=np.float32)
+        new_precisions = np.interp(new_thresholds, thresholds, precisions)
+        new_recalls = np.interp(new_thresholds, thresholds, recalls)
+        thresholds = new_thresholds
+        precisions = new_precisions
+        recalls = new_recalls
+
+    return precisions.astype(np.float32), recalls.astype(np.float32), thresholds.astype(np.float32)
+
+
+def calculate_roc_curve(
+    references: NDArray[np.int64],
+    probabilities: NDArray[np.float32],
+    max_length: int = 100,
+) -> Tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:
+    if probabilities.ndim == 1:
+        probabilities_slice = probabilities
+    elif probabilities.ndim == 2:
+        probabilities_slice = probabilities[:, 1]
+    else:
+        raise ValueError("Probabilities must be 1 or 2 dimensional")
+
+    if len(probabilities_slice) != len(references):
+        raise ValueError("Probabilities and references must have the same length")
+
+    # Convert probabilities to float32 before calling sklearn_roc_curve
+    probabilities_slice = probabilities_slice.astype(np.float32)
+    fpr, tpr, thresholds = sklearn_roc_curve(references, probabilities_slice)
+
+    # Convert all arrays to float32 immediately after getting them
+    fpr = fpr.astype(np.float32)
+    tpr = tpr.astype(np.float32)
+    thresholds = thresholds.astype(np.float32)
+
+    # We set the first threshold to 1.0 instead of inf for reasonable values in interpolation
+    thresholds[0] = 1.0
+
+    # Sort by threshold
+    sorted_indices = np.argsort(thresholds)
+    thresholds = thresholds[sorted_indices]
+    fpr = fpr[sorted_indices]
+    tpr = tpr[sorted_indices]
+
+    if len(fpr) > max_length:
+        new_thresholds = np.linspace(0, 1, max_length, dtype=np.float32)
+        new_fpr = np.interp(new_thresholds, thresholds, fpr)
+        new_tpr = np.interp(new_thresholds, thresholds, tpr)
+        thresholds = new_thresholds
+        fpr = new_fpr
+        tpr = new_tpr
+
+    return fpr.astype(np.float32), tpr.astype(np.float32), thresholds.astype(np.float32)

orca_sdk/_shared/metrics_test.py

@@ -0,0 +1,169 @@
+"""
+IMPORTANT:
+- This is a shared file between OrcaLib and the Orca SDK.
+- Please ensure that it does not have any dependencies on the OrcaLib code.
+- Make sure to edit this file in orcalib/shared and NOT in orca_sdk, since it will be overwritten there.
+"""
+
+from typing import Literal
+
+import numpy as np
+import pytest
+
+from .metrics import (
+    EvalPrediction,
+    calculate_pr_curve,
+    calculate_roc_curve,
+    classification_scores,
+    compute_classifier_metrics,
+    softmax,
+)
+
+
+def test_binary_metrics():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.3, 0.2])
+
+    metrics = classification_scores(y_true, y_score)
+
+    assert metrics["accuracy"] == 0.8
+    assert metrics["f1_score"] == 0.8
+    assert metrics["roc_auc"] is not None
+    assert metrics["roc_auc"] > 0.8
+    assert metrics["roc_auc"] < 1.0
+    assert metrics["pr_auc"] is not None
+    assert metrics["pr_auc"] > 0.8
+    assert metrics["pr_auc"] < 1.0
+    assert metrics["log_loss"] is not None
+    assert metrics["log_loss"] > 0.0
+
+
+def test_multiclass_metrics_with_2_classes():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([[0.9, 0.1], [0.1, 0.9], [0.2, 0.8], [0.7, 0.3], [0.8, 0.2]])
+
+    metrics = classification_scores(y_true, y_score)
+
+    assert metrics["accuracy"] == 0.8
+    assert metrics["f1_score"] == 0.8
+    assert metrics["roc_auc"] is not None
+    assert metrics["roc_auc"] > 0.8
+    assert metrics["roc_auc"] < 1.0
+    assert metrics["pr_auc"] is not None
+    assert metrics["pr_auc"] > 0.8
+    assert metrics["pr_auc"] < 1.0
+    assert metrics["log_loss"] is not None
+    assert metrics["log_loss"] > 0.0
+
+
+@pytest.mark.parametrize(
+    "average, multiclass",
+    [("micro", "ovr"), ("macro", "ovr"), ("weighted", "ovr"), ("micro", "ovo"), ("macro", "ovo"), ("weighted", "ovo")],
+)
+def test_multiclass_metrics_with_3_classes(
+    average: Literal["micro", "macro", "weighted"], multiclass: Literal["ovr", "ovo"]
+):
+    y_true = np.array([0, 1, 1, 0, 2])
+    y_score = np.array([[0.9, 0.1, 0.0], [0.1, 0.9, 0.0], [0.2, 0.8, 0.0], [0.7, 0.3, 0.0], [0.0, 0.0, 1.0]])
+
+    metrics = classification_scores(y_true, y_score, average=average, multi_class=multiclass)
+
+    assert metrics["accuracy"] == 1.0
+    assert metrics["f1_score"] == 1.0
+    assert metrics["roc_auc"] is not None
+    assert metrics["roc_auc"] > 0.8
+    assert metrics["pr_auc"] is None
+    assert metrics["log_loss"] is not None
+    assert metrics["log_loss"] > 0.0
+
+
+def test_does_not_modify_logits_unless_necessary():
+    logits = np.array([[0.1, 0.9], [0.2, 0.8], [0.7, 0.3], [0.8, 0.2]])
+    references = np.array([0, 1, 0, 1])
+    metrics = compute_classifier_metrics(EvalPrediction(logits, references))
+    assert metrics["log_loss"] == classification_scores(references, logits)["log_loss"]
+
+
+def test_normalizes_logits_if_necessary():
+    logits = np.array([[1.2, 3.9], [1.2, 5.8], [1.2, 2.7], [1.2, 1.3]])
+    references = np.array([0, 1, 0, 1])
+    metrics = compute_classifier_metrics(EvalPrediction(logits, references))
+    assert (
+        metrics["log_loss"] == classification_scores(references, logits / logits.sum(axis=1, keepdims=True))["log_loss"]
+    )
+
+
+def test_softmaxes_logits_if_necessary():
+    logits = np.array([[-1.2, 3.9], [1.2, -5.8], [1.2, 2.7], [1.2, 1.3]])
+    references = np.array([0, 1, 0, 1])
+    metrics = compute_classifier_metrics(EvalPrediction(logits, references))
+    assert metrics["log_loss"] == classification_scores(references, softmax(logits))["log_loss"]
+
+
+def test_precision_recall_curve():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    precision, recall, thresholds = calculate_pr_curve(y_true, y_score)
+    assert precision is not None
+    assert recall is not None
+    assert thresholds is not None
+
+    assert len(precision) == len(recall) == len(thresholds) == 6
+    assert precision[0] == 0.6
+    assert recall[0] == 1.0
+    assert precision[-1] == 1.0
+    assert recall[-1] == 0.0
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(thresholds) >= 0)
+
+
+def test_roc_curve():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    fpr, tpr, thresholds = calculate_roc_curve(y_true, y_score)
+    assert fpr is not None
+    assert tpr is not None
+    assert thresholds is not None
+
+    assert len(fpr) == len(tpr) == len(thresholds) == 6
+    assert fpr[0] == 1.0
+    assert tpr[0] == 1.0
+    assert fpr[-1] == 0.0
+    assert tpr[-1] == 0.0
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(thresholds) >= 0)
+
+
+def test_precision_recall_curve_max_length():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    precision, recall, thresholds = calculate_pr_curve(y_true, y_score, max_length=5)
+    assert len(precision) == len(recall) == len(thresholds) == 5
+
+    assert precision[0] == 0.6
+    assert recall[0] == 1.0
+    assert precision[-1] == 1.0
+    assert recall[-1] == 0.0
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(thresholds) >= 0)
+
+
+def test_roc_curve_max_length():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    fpr, tpr, thresholds = calculate_roc_curve(y_true, y_score, max_length=5)
+    assert len(fpr) == len(tpr) == len(thresholds) == 5
+    assert fpr[0] == 1.0
+    assert tpr[0] == 1.0
+    assert fpr[-1] == 0.0
+    assert tpr[-1] == 0.0
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(thresholds) >= 0)

orca_sdk/classification_model.py

@@ -6,6 +6,15 @@ from datetime import datetime
 from typing import Any, Generator, Iterable, Literal, cast, overload
 from uuid import UUID
 
+import numpy as np
+from datasets import Dataset
+from sklearn.metrics import (
+    accuracy_score,
+    auc,
+    f1_score,
+    roc_auc_score,
+)
+
 from ._generated_api_client.api import (
     create_evaluation,
     create_model,
@@ -19,9 +28,11 @@ from ._generated_api_client.api import (
     update_model,
 )
 from ._generated_api_client.models import (
+    ClassificationEvaluationResult,
     CreateRACModelRequest,
     EvaluationRequest,
     ListPredictionsRequest,
+    PrecisionRecallCurve,
 )
 from ._generated_api_client.models import (
     PredictionSortItemItemType0 as PredictionSortColumns,
@@ -33,8 +44,10 @@ from ._generated_api_client.models import (
     RACHeadType,
     RACModelMetadata,
     RACModelUpdate,
+    ROCCurve,
 )
 from ._generated_api_client.models.prediction_request import PredictionRequest
+from ._shared.metrics import calculate_pr_curve, calculate_roc_curve
 from ._utils.common import UNSET, CreateMode, DropMode
 from ._utils.task import wait_for_task
 from .datasource import Datasource
@@ -372,6 +385,7 @@ class ClassificationModel:
                 anomaly_score=prediction.anomaly_score,
                 memoryset=self.memoryset,
                 model=self,
+                logits=prediction.logits,
             )
             for prediction in response
         ]
@@ -444,46 +458,157 @@ class ClassificationModel:
             for prediction in predictions
         ]
 
-    def evaluate(
+    def _calculate_metrics(
+        self,
+        predictions: list[LabelPrediction],
+        expected_labels: list[int],
+    ) -> ClassificationEvaluationResult:
+
+        targets_array = np.array(expected_labels)
+        predictions_array = np.array([p.label for p in predictions])
+
+        logits_array = np.array([p.logits for p in predictions])
+
+        f1 = float(f1_score(targets_array, predictions_array, average="weighted"))
+        accuracy = float(accuracy_score(targets_array, predictions_array))
+
+        # Only compute ROC AUC and PR AUC for binary classification
+        unique_classes = np.unique(targets_array)
+
+        pr_curve = None
+        roc_curve = None
+
+        if len(unique_classes) == 2:
+            try:
+                precisions, recalls, pr_thresholds = calculate_pr_curve(targets_array, logits_array)
+                pr_auc = float(auc(recalls, precisions))
+
+                pr_curve = PrecisionRecallCurve(
+                    precisions=precisions.tolist(),
+                    recalls=recalls.tolist(),
+                    thresholds=pr_thresholds.tolist(),
+                    auc=pr_auc,
+                )
+
+                fpr, tpr, roc_thresholds = calculate_roc_curve(targets_array, logits_array)
+                roc_auc = float(roc_auc_score(targets_array, logits_array[:, 1]))
+
+                roc_curve = ROCCurve(
+                    false_positive_rates=fpr.tolist(),
+                    true_positive_rates=tpr.tolist(),
+                    thresholds=roc_thresholds.tolist(),
+                    auc=roc_auc,
+                )
+            except ValueError as e:
+                logging.warning(f"Error calculating PR and ROC curves: {e}")
+
+        return ClassificationEvaluationResult(
+            f1_score=f1,
+            accuracy=accuracy,
+            loss=0.0,
+            precision_recall_curve=pr_curve,
+            roc_curve=roc_curve,
+        )
+
+    def _evaluate_datasource(
         self,
         datasource: Datasource,
+        value_column: str,
+        label_column: str,
+        record_predictions: bool,
+        tags: set[str] | None,
+    ) -> dict[str, Any]:
+        response = create_evaluation(
+            self.id,
+            body=EvaluationRequest(
+                datasource_id=datasource.id,
+                datasource_label_column=label_column,
+                datasource_value_column=value_column,
+                memoryset_override_id=self._memoryset_override_id,
+                record_telemetry=record_predictions,
+                telemetry_tags=list(tags) if tags else None,
+            ),
+        )
+        wait_for_task(response.task_id, description="Running evaluation")
+        response = get_evaluation(self.id, UUID(response.task_id))
+        assert response.result is not None
+        return response.result.to_dict()
+
+    def _evaluate_dataset(
+        self,
+        dataset: Dataset,
+        value_column: str,
+        label_column: str,
+        record_predictions: bool,
+        tags: set[str],
+        batch_size: int,
+    ) -> dict[str, Any]:
+        predictions = []
+        expected_labels = []
+
+        for i in range(0, len(dataset), batch_size):
+            batch = dataset[i : i + batch_size]
+            predictions.extend(
+                self.predict(
+                    batch[value_column],
+                    expected_labels=batch[label_column],
+                    tags=tags,
+                    disable_telemetry=(not record_predictions),
+                )
+            )
+            expected_labels.extend(batch[label_column])
+
+        return self._calculate_metrics(predictions, expected_labels).to_dict()
+
+    def evaluate(
+        self,
+        data: Datasource | Dataset,
         value_column: str = "value",
         label_column: str = "label",
         record_predictions: bool = False,
-        tags: set[str] | None = None,
+        tags: set[str] = {"evaluation"},
+        batch_size: int = 100,
     ) -> dict[str, Any]:
         """
-        Evaluate the classification model on a given datasource
+        Evaluate the classification model on a given dataset or datasource
 
         Params:
-            datasource: Datasource to evaluate the model on
+            data: Dataset or Datasource to evaluate the model on
             value_column: Name of the column that contains the input values to the model
             label_column: Name of the column containing the expected labels
             record_predictions: Whether to record [`LabelPrediction`][orca_sdk.telemetry.LabelPrediction]s for analysis
             tags: Optional tags to add to the recorded [`LabelPrediction`][orca_sdk.telemetry.LabelPrediction]s
+            batch_size: Batch size for processing Dataset inputs (only used when input is a Dataset)
 
         Returns:
             Dictionary with evaluation metrics
 
         Examples:
+            Evaluate using a Datasource:
             >>> model.evaluate(datasource, value_column="text", label_column="airline_sentiment")
             { "f1_score": 0.85, "roc_auc": 0.85, "pr_auc": 0.85, "accuracy": 0.85, "loss": 0.35, ... }
+
+            Evaluate using a Dataset:
+            >>> model.evaluate(dataset, value_column="text", label_column="sentiment")
+            { "f1_score": 0.85, "roc_auc": 0.85, "pr_auc": 0.85, "accuracy": 0.85, "loss": 0.35, ... }
         """
-        response = create_evaluation(
-            self.id,
-            body=EvaluationRequest(
-                datasource_id=datasource.id,
-                datasource_label_column=label_column,
-                datasource_value_column=value_column,
-                memoryset_override_id=self._memoryset_override_id,
-                record_telemetry=record_predictions,
-                telemetry_tags=list(tags) if tags else None,
-            ),
-        )
-        wait_for_task(response.task_id, description="Running evaluation")
-        response = get_evaluation(self.id, UUID(response.task_id))
-        assert response.result is not None
-        return response.result.to_dict()
+        if isinstance(data, Datasource):
+            return self._evaluate_datasource(
+                datasource=data,
+                value_column=value_column,
+                label_column=label_column,
+                record_predictions=record_predictions,
+                tags=tags,
+            )
+        else:
+            return self._evaluate_dataset(
+                dataset=data,
+                value_column=value_column,
+                label_column=label_column,
+                record_predictions=record_predictions,
+                tags=tags,
+                batch_size=batch_size,
+            )
 
     def finetune(self, datasource: Datasource):
         #  do not document until implemented

orca_sdk/classification_model_test.py

@@ -1,5 +1,6 @@
 from uuid import uuid4
 
+import numpy as np
 import pytest
 from datasets.arrow_dataset import Dataset
 
@@ -138,28 +139,47 @@ def test_delete_memoryset_before_model_constraint_violation(hf_dataset):
         LabeledMemoryset.drop(memoryset.id)
 
 
-def test_evaluate(model):
-    eval_datasource = Datasource.from_list(
-        "eval_datasource",
-        [
-            {"text": "chicken noodle soup is the best", "label": 1},
-            {"text": "cats are cute", "label": 0},
-            {"text": "soup is great for the winter", "label": 0},
-            {"text": "i love cats", "label": 1},
-        ],
-    )
-    result = model.evaluate(eval_datasource, value_column="text")
-    assert result is not None
-    assert isinstance(result, dict)
-    assert isinstance(result["accuracy"], float)
-    assert isinstance(result["f1_score"], float)
-    assert isinstance(result["loss"], float)
-    assert len(result["precision_recall_curve"]["thresholds"]) == 4
-    assert len(result["precision_recall_curve"]["precisions"]) == 4
-    assert len(result["precision_recall_curve"]["recalls"]) == 4
-    assert len(result["roc_curve"]["thresholds"]) == 4
-    assert len(result["roc_curve"]["false_positive_rates"]) == 4
-    assert len(result["roc_curve"]["true_positive_rates"]) == 4
+def test_evaluate_combined(model):
+    data = [
+        {"text": "chicken noodle soup is the best", "label": 1},
+        {"text": "cats are cute", "label": 0},
+        {"text": "soup is great for the winter", "label": 0},
+        {"text": "i love cats", "label": 1},
+    ]
+
+    eval_datasource = Datasource.from_list("eval_datasource", data)
+    result_datasource = model.evaluate(eval_datasource, value_column="text")
+
+    eval_dataset = Dataset.from_list(data)
+    result_dataset = model.evaluate(eval_dataset, value_column="text")
+
+    for result in [result_datasource, result_dataset]:
+        assert result is not None
+        assert isinstance(result, dict)
+        assert isinstance(result["accuracy"], float)
+        assert isinstance(result["f1_score"], float)
+        assert isinstance(result["loss"], float)
+        assert np.allclose(result["accuracy"], 0.5)
+        assert np.allclose(result["f1_score"], 0.5)
+
+        assert isinstance(result["precision_recall_curve"]["thresholds"], list)
+        assert isinstance(result["precision_recall_curve"]["precisions"], list)
+        assert isinstance(result["precision_recall_curve"]["recalls"], list)
+        assert isinstance(result["roc_curve"]["thresholds"], list)
+        assert isinstance(result["roc_curve"]["false_positive_rates"], list)
+        assert isinstance(result["roc_curve"]["true_positive_rates"], list)
+
+        assert np.allclose(result["roc_curve"]["thresholds"], [0.0, 0.8155114054679871, 0.834095299243927, 1.0])
+        assert np.allclose(result["roc_curve"]["false_positive_rates"], [1.0, 0.5, 0.0, 0.0])
+        assert np.allclose(result["roc_curve"]["true_positive_rates"], [1.0, 0.5, 0.5, 0.0])
+        assert np.allclose(result["roc_curve"]["auc"], 0.625)
+
+        assert np.allclose(
+            result["precision_recall_curve"]["thresholds"], [0.0, 0.0, 0.8155114054679871, 0.834095299243927]
+        )
+        assert np.allclose(result["precision_recall_curve"]["precisions"], [0.5, 0.5, 1.0, 1.0])
+        assert np.allclose(result["precision_recall_curve"]["recalls"], [1.0, 0.5, 0.5, 0.0])
+        assert np.allclose(result["precision_recall_curve"]["auc"], 0.75)
 
 
 def test_evaluate_with_telemetry(model):
@@ -188,6 +208,13 @@ def test_predict(model: ClassificationModel, label_names: list[str]):
     assert predictions[1].label_name == label_names[1]
     assert 0 <= predictions[1].confidence <= 1
 
+    assert predictions[0].logits is not None
+    assert predictions[1].logits is not None
+    assert len(predictions[0].logits) == 2
+    assert len(predictions[1].logits) == 2
+    assert predictions[0].logits[0] > predictions[0].logits[1]
+    assert predictions[1].logits[0] < predictions[1].logits[1]
+
 
 def test_predict_disable_telemetry(model: ClassificationModel, label_names: list[str]):
     predictions = model.predict(["Do you love soup?", "Are cats cute?"], disable_telemetry=True)

orca_sdk/telemetry.py

@@ -135,6 +135,7 @@ class LabelPrediction:
     anomaly_score: float | None
     memoryset: LabeledMemoryset
     model: ClassificationModel
+    logits: list[float] | None
 
     def __init__(
         self,
@@ -147,6 +148,7 @@ class LabelPrediction:
         memoryset: LabeledMemoryset | str,
         model: ClassificationModel | str,
         telemetry: LabelPredictionWithMemoriesAndFeedback | None = None,
+        logits: list[float] | None = None,
     ):
         # for internal use only, do not document
         from .classification_model import ClassificationModel
@@ -159,6 +161,7 @@ class LabelPrediction:
         self.memoryset = LabeledMemoryset.open(memoryset) if isinstance(memoryset, str) else memoryset
         self.model = ClassificationModel.open(model) if isinstance(model, str) else model
         self.__telemetry = telemetry if telemetry else None
+        self.logits = logits
 
     def __repr__(self):
         return (

{orca_sdk-0.0.91.dist-info → orca_sdk-0.0.92.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: orca_sdk
-Version: 0.0.91
+Version: 0.0.92
 Summary: SDK for interacting with Orca Services
 License: Apache-2.0
 Author: Orca DB Inc.
@@ -20,7 +20,9 @@ Requires-Dist: pandas (>=2.2.3,<3.0.0)
 Requires-Dist: pyarrow (>=18.0.0,<19.0.0)
 Requires-Dist: python-dateutil (>=2.8.0,<3.0.0)
 Requires-Dist: python-dotenv (>=1.1.0,<2.0.0)
+Requires-Dist: scikit-learn (>=1.6.1,<2.0.0)
 Requires-Dist: torch (>=2.5.1,<3.0.0)
+Requires-Dist: transformers (>=4.51.3,<5.0.0)
 Description-Content-Type: text/markdown
 
 <!--

{orca_sdk-0.0.91.dist-info → orca_sdk-0.0.92.dist-info}/RECORD

@@ -81,7 +81,7 @@ orca_sdk/_generated_api_client/models/__init__.py,sha256=3fjbYdRtS5POw4Ce2FfBdnU
 orca_sdk/_generated_api_client/models/analyze_neighbor_labels_result.py,sha256=n5xhKNRx_QaAmXgveWSwLRlAjTHkuEGiH0-Vr1H6RsY,4256
 orca_sdk/_generated_api_client/models/api_key_metadata.py,sha256=jQrSe_X5hCgFYh8PwX-X0M6VINVGVhLBlKmv4qN5otA,3789
 orca_sdk/_generated_api_client/models/api_key_metadata_scope_item.py,sha256=umhWzrAt0ZEI9B7pLpnOEsc5Kc-dxeJdzHX7iHyjt4I,182
-orca_sdk/_generated_api_client/models/base_label_prediction_result.py,sha256=adYVUJR_aFApkmUNOCnhBxtRmhk6revQAsXPsoED-UY,3624
+orca_sdk/_generated_api_client/models/base_label_prediction_result.py,sha256=wJBkJcUdI588tOXimOZ6lBIFGPAaStBrOC84m4-8CIw,3828
 orca_sdk/_generated_api_client/models/base_model.py,sha256=0UY9I_q-b6kOG0LYcw_C192PKRfmejYX9rZa7POCrTc,1563
 orca_sdk/_generated_api_client/models/body_create_datasource_datasource_post.py,sha256=w5Ni3zFPMTf8XYsH-EZmlokx7vV0vpQiSnbXlktoQBY,6713
 orca_sdk/_generated_api_client/models/classification_evaluation_result.py,sha256=mdSZjv7qy6OreEjwNTV_VpfoeuZHdrnlCG8sr0elhoo,4715
@@ -121,7 +121,7 @@ orca_sdk/_generated_api_client/models/internal_server_error_response.py,sha256=R
 orca_sdk/_generated_api_client/models/label_class_metrics.py,sha256=Q3vWLw8F_IdwAwhunLp0f_l7PvP1gZN1XGCZQRJtbAY,3144
 orca_sdk/_generated_api_client/models/label_prediction_memory_lookup.py,sha256=DEwdX5532kHRpsKJe3wOgUWUTZOdeaJV30XvsI8dyOI,6005
 orca_sdk/_generated_api_client/models/label_prediction_memory_lookup_metadata.py,sha256=bsXtXEf25ch5qAdpnXWSi2qzCkQPZ4xhKcHWMxlgOhQ,2338
-orca_sdk/_generated_api_client/models/label_prediction_with_memories_and_feedback.py,sha256=fa0djMjPuJkEwRlRvmLiVHEJHraNfNxoANC0ACwC-ZQ,9004
+orca_sdk/_generated_api_client/models/label_prediction_with_memories_and_feedback.py,sha256=ZxmUTIY02_eZz5EeO6xwehFyS4yYzt_Nw3v0pEbTclQ,9004
 orca_sdk/_generated_api_client/models/labeled_memory.py,sha256=BYG1PqvL3FXKQCuBTg3pLwIgA0Uv8KU5YoxvdR2zZxg,5205
 orca_sdk/_generated_api_client/models/labeled_memory_insert.py,sha256=O3rgrloH3eu9YPzP7X1AKRfq6wxx9Eznl_prpRiMVVM,3768
 orca_sdk/_generated_api_client/models/labeled_memory_insert_metadata.py,sha256=b9T8i61YUIkNpbJzLwUztfUGNqwRzYJ51RfFpukNS5I,2295
@@ -134,7 +134,7 @@ orca_sdk/_generated_api_client/models/labeled_memory_update_metadata_type_0.py,s
 orca_sdk/_generated_api_client/models/labeled_memory_with_feedback_metrics.py,sha256=pr2cM9z2F4iAMW6N38xGCYD_fr8R5co70-p0TVRN94w,6307
 orca_sdk/_generated_api_client/models/labeled_memory_with_feedback_metrics_feedback_metrics.py,sha256=qSUl04WibCHV-1yoytEW2TI5in2cf1HCerpOJ8wej3w,2272
 orca_sdk/_generated_api_client/models/labeled_memory_with_feedback_metrics_metadata.py,sha256=wYnU5KuMTlUwIxpbrCe4obx40h_-FJExxoCOMd0-Qik,2366
-orca_sdk/_generated_api_client/models/labeled_memoryset_metadata.py,sha256=x-o5h1zymo9Wn-ae4Wh1d1eafDs7PUcN4kjztGA6lXg,7326
+orca_sdk/_generated_api_client/models/labeled_memoryset_metadata.py,sha256=1GiEJSXBA6VnUoLYSNk2f5Zxyj8bi_kWvUGOUzHNKyQ,7669
 orca_sdk/_generated_api_client/models/labeled_memoryset_update.py,sha256=xd5obMpcK1zZiU-q4xQUbrWnkaIi176gcjZTBoAMlpQ,3586
 orca_sdk/_generated_api_client/models/list_memories_request.py,sha256=ZPp2FR8-tNMc9eAmErAHEpLf2xrvI_6NtGldSQfAfe4,3091
 orca_sdk/_generated_api_client/models/list_predictions_request.py,sha256=I20mJhJhx-sIeFeK1WNbmaTI07U2lhS840pURBZdYGo,9976
@@ -190,6 +190,9 @@ orca_sdk/_generated_api_client/models/unauthorized_error_response.py,sha256=Sr-p
 orca_sdk/_generated_api_client/models/update_prediction_request.py,sha256=HMPq_K0MlQY7beWn73LEhgjUNcBEjqGC8oFlB9t9em0,3573
 orca_sdk/_generated_api_client/py.typed,sha256=8ZJUsxZiuOy1oJeVhsTWQhTG_6pTVHVXk5hJL79ebTk,25
 orca_sdk/_generated_api_client/types.py,sha256=j7-uA7wWwN1cq0d7ULccN4vDm-1IzgnrxSyVktxvABM,1399
+orca_sdk/_shared/__init__.py,sha256=aXGbM6K8IN5V_7bPeTQZE2CZedV1i1IkynS7swq8D7k,89
+orca_sdk/_shared/metrics.py,sha256=FNGOSfZke3AVCf-j7FdYcq7nmH68RJ0SqD0r_LsLjeY,7565
+orca_sdk/_shared/metrics_test.py,sha256=Udv_JsHbYFYtP2W7iFHgnafOciD03te25qvrX9PUaQ8,5522
 orca_sdk/_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 orca_sdk/_utils/analysis_ui.py,sha256=2ADUOxbLTcY0sYOpHeReTr13sQ7Yu4kZQ66RDeUuWZw,9216
 orca_sdk/_utils/analysis_ui_style.css,sha256=q_ba_-_KtgztepHg829zLzypaxKayl7ySC1-oYDzV3k,836
@@ -203,8 +206,8 @@ orca_sdk/_utils/prediction_result_ui.py,sha256=dudc21ka2Bqdtr_8wQaMVFxLGvrsZxWUZ
 orca_sdk/_utils/task.py,sha256=WOfFuRCoh6QHVDDYrGeq3Hi6NVihQQZJii0cBCONOWk,2400
 orca_sdk/_utils/value_parser.py,sha256=c3qMABCCDQcIjn9N1orYYnlRwDW9JWdGwW_2TDZPLdI,1286
 orca_sdk/_utils/value_parser_test.py,sha256=OybsiC-Obi32RRi9NIuwrVBRAnlyPMV1xVAaevSrb7M,1079
-orca_sdk/classification_model.py,sha256=08zEj4HCn5dws7s05w8I4PUBcE76RySew5P9nIM7sHE,22143
-orca_sdk/classification_model_test.py,sha256=48Zcl4k6oIPYBuTTG6rdrEigijdFA-7wRFo3tHG4Osc,11811
+orca_sdk/classification_model.py,sha256=j3b277NGeF2kDehwPN7s95KkEyjafvH52ip9t-dRFPk,26439
+orca_sdk/classification_model_test.py,sha256=gbqyjjnwVZB_Z7IHLIVJ7U3jY-xdKag9hIXULDelVqQ,13272
 orca_sdk/conftest.py,sha256=_7O6yVccU-_zteUTCX3j7j7ZfyKNBD7nYL-G8ln6qXY,4661
 orca_sdk/credentials.py,sha256=gq_4w_o-igCCLNR6TY1x4RzMYysKUCsXJvdi6nem-A0,3558
 orca_sdk/credentials_test.py,sha256=ETTyDZ9MEpb_X6yiRcgYGWNKCB2QZ5CLYB_unRGg1b8,1028
@@ -214,8 +217,8 @@ orca_sdk/embedding_model.py,sha256=Hw8NlwzWVK5ts8SF0lHIs7hL38hCTreEiIyoqHY-OFA,1
 orca_sdk/embedding_model_test.py,sha256=j6uGu9ZJSafDV7uFiJiG8SZVGvPQBgxxDcg7i1xbWho,6914
 orca_sdk/memoryset.py,sha256=xvaNn3YwG3fzk3MZhk3LeX_K5yRKP-yRf79bIAUBR-Y,56058
 orca_sdk/memoryset_test.py,sha256=w8-2RXFePg1pqC67uMpHSevjnW4P0GbNpqRjJXAmIa0,15122
-orca_sdk/telemetry.py,sha256=RqVL8Q7-NMCSlygFpetvmnFDDmfrG1T4GuPA7llCIN8,16290
+orca_sdk/telemetry.py,sha256=U53NI7_D1IpWqdV8NYuUrwvhpX0CF_PJvRRvOiFekno,16393
 orca_sdk/telemetry_test.py,sha256=7JfS0k7r9STMCkasCjXWL3KmbrdmVjVnFeYPCdT8jqQ,5059
-orca_sdk-0.0.91.dist-info/METADATA,sha256=FwJ_e6DFRNWsrMijiSeEmuuvzHNXRqvoBBUYT6yYHn4,3138
-orca_sdk-0.0.91.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-orca_sdk-0.0.91.dist-info/RECORD,,
+orca_sdk-0.0.92.dist-info/METADATA,sha256=W1Ee5mRfJQH03lS5XQAztQC1cE4gupmRJaibugGAcrs,3229
+orca_sdk-0.0.92.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+orca_sdk-0.0.92.dist-info/RECORD,,