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

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/_utils/analysis_ui.py

@@ -152,7 +152,7 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
                     predicted_label_name = label_names[predicted_label]
                     predicted_label_confidence = mem.metrics.get("neighbor_predicted_label_confidence", 0)
 
-                    with gr.Row(equal_height=True, variant="panel", elem_classes="white" if i % 2 == 0 else None):
+                    with gr.Row(equal_height=True, variant="panel"):
                         with gr.Column(scale=9):
                             assert isinstance(mem.value, str)
                             gr.Markdown(mem.value, label="Value", height=50)

orca_sdk/_utils/analysis_ui_style.css

@@ -1,6 +1,3 @@
-.white {
-  background-color: white;
-}
 .centered input {
   margin: auto;
 }

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,
@@ -16,11 +25,14 @@ from ._generated_api_client.api import (
     list_predictions,
     predict_gpu,
     record_prediction_feedback,
+    update_model,
 )
 from ._generated_api_client.models import (
+    ClassificationEvaluationResult,
     CreateRACModelRequest,
     EvaluationRequest,
     ListPredictionsRequest,
+    PrecisionRecallCurve,
 )
 from ._generated_api_client.models import (
     PredictionSortItemItemType0 as PredictionSortColumns,
@@ -31,9 +43,12 @@ from ._generated_api_client.models import (
 from ._generated_api_client.models import (
     RACHeadType,
     RACModelMetadata,
+    RACModelUpdate,
+    ROCCurve,
 )
 from ._generated_api_client.models.prediction_request import PredictionRequest
-from ._utils.common import CreateMode, DropMode
+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
 from .memoryset import LabeledMemoryset
@@ -270,18 +285,53 @@ class ClassificationModel:
             if if_not_exists == "error":
                 raise
 
+    def refresh(self):
+        """Refresh the model data from the OrcaCloud"""
+        self.__dict__.update(ClassificationModel.open(self.name).__dict__)
+
+    def update_metadata(self, *, description: str | None = UNSET) -> None:
+        """
+        Update editable classification model metadata properties.
+
+        Params:
+            description: Value to set for the description, defaults to `[UNSET]` if not provided.
+
+        Examples:
+            Update the description:
+            >>> model.update(description="New description")
+
+            Remove description:
+            >>> model.update(description=None)
+        """
+        update_model(self.id, body=RACModelUpdate(description=description))
+        self.refresh()
+
     @overload
     def predict(
-        self, value: list[str], expected_labels: list[int] | None = None, tags: set[str] = set()
+        self,
+        value: list[str],
+        expected_labels: list[int] | None = None,
+        tags: set[str] = set(),
+        disable_telemetry: bool = False,
     ) -> list[LabelPrediction]:
         pass
 
     @overload
-    def predict(self, value: str, expected_labels: int | None = None, tags: set[str] = set()) -> LabelPrediction:
+    def predict(
+        self,
+        value: str,
+        expected_labels: int | None = None,
+        tags: set[str] = set(),
+        disable_telemetry: bool = False,
+    ) -> LabelPrediction:
         pass
 
     def predict(
-        self, value: list[str] | str, expected_labels: list[int] | int | None = None, tags: set[str] = set()
+        self,
+        value: list[str] | str,
+        expected_labels: list[int] | int | None = None,
+        tags: set[str] = set(),
+        disable_telemetry: bool = False,
     ) -> list[LabelPrediction] | LabelPrediction:
         """
         Predict label(s) for the given input value(s) grounded in similar memories
@@ -290,6 +340,7 @@ class ClassificationModel:
             value: Value(s) to get predict the labels of
             expected_labels: Expected label(s) for the given input to record for model evaluation
             tags: Tags to add to the prediction(s)
+            disable_telemetry: Whether to disable telemetry for the prediction(s)
 
         Returns:
             Label prediction or list of label predictions
@@ -318,8 +369,13 @@ class ClassificationModel:
                     else [expected_labels] if expected_labels is not None else None
                 ),
                 tags=list(tags),
+                disable_telemetry=disable_telemetry,
             ),
         )
+
+        if not disable_telemetry and any(p.prediction_id is None for p in response):
+            raise RuntimeError("Failed to save prediction to database.")
+
         predictions = [
             LabelPrediction(
                 prediction_id=prediction.prediction_id,
@@ -329,6 +385,7 @@ class ClassificationModel:
                 anomaly_score=prediction.anomaly_score,
                 memoryset=self.memoryset,
                 model=self,
+                logits=prediction.logits,
             )
             for prediction in response
         ]
@@ -401,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