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

orca_sdk/classification_model.py

@@ -15,6 +15,7 @@ from .client import (
     BootstrapLabeledMemoryDataResult,
     ClassificationModelMetadata,
     ClassificationPredictionRequest,
+    ListPredictionsRequest,
     OrcaClient,
     PredictiveModelUpdate,
     RACHeadType,
@@ -363,6 +364,7 @@ class ClassificationModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
         use_gpu: bool = True,
+        batch_size: int = 100,
     ) -> list[ClassificationPrediction]:
         pass
 
@@ -383,6 +385,7 @@ class ClassificationModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
         use_gpu: bool = True,
+        batch_size: int = 100,
     ) -> ClassificationPrediction:
         pass
 
@@ -402,6 +405,7 @@ class ClassificationModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
         use_gpu: bool = True,
+        batch_size: int = 100,
     ) -> list[ClassificationPrediction] | ClassificationPrediction:
         """
         Predict label(s) for the given input value(s) grounded in similar memories
@@ -429,6 +433,7 @@ class ClassificationModel:
                 * `"exclude_global"`: Exclude global memories
                 * `"only_global"`: Only include global memories
             use_gpu: Whether to use GPU for the prediction (defaults to True)
+            batch_size: Number of values to process in a single API call
 
         Returns:
             Label prediction or list of label predictions
@@ -456,6 +461,8 @@ class ClassificationModel:
 
         if timeout_seconds <= 0:
             raise ValueError("timeout_seconds must be a positive integer")
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
 
         parsed_filters = [
             _parse_filter_item_from_tuple(filter) if isinstance(filter, tuple) else filter for filter in filters
@@ -464,10 +471,17 @@ class ClassificationModel:
         if any(_is_metric_column(filter[0]) for filter in filters):
             raise ValueError(f"Cannot filter on {filters} - telemetry filters are not supported for predictions")
 
+        # Convert to list for batching
+        values = value if isinstance(value, list) else [value]
+        if isinstance(expected_labels, list) and len(expected_labels) != len(values):
+            raise ValueError("Invalid input: \n\texpected_labels must be the same length as values")
+        if isinstance(partition_id, list) and len(partition_id) != len(values):
+            raise ValueError("Invalid input: \n\tpartition_id must be the same length as values")
+
         if isinstance(expected_labels, int):
-            expected_labels = [expected_labels]
+            expected_labels = [expected_labels] * len(values)
         elif isinstance(expected_labels, str):
-            expected_labels = [self.memoryset.label_names.index(expected_labels)]
+            expected_labels = [self.memoryset.label_names.index(expected_labels)] * len(values)
         elif isinstance(expected_labels, list):
             expected_labels = [
                 self.memoryset.label_names.index(label) if isinstance(label, str) else label
@@ -481,47 +495,56 @@ class ClassificationModel:
 
         telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
         client = OrcaClient._resolve_client()
-        request_json: ClassificationPredictionRequest = {
-            "input_values": value if isinstance(value, list) else [value],
-            "memoryset_override_name_or_id": self._memoryset_override_id,
-            "expected_labels": expected_labels,
-            "tags": list(tags or set()),
-            "save_telemetry": telemetry_on,
-            "save_telemetry_synchronously": telemetry_sync,
-            "filters": cast(list[FilterItem], parsed_filters),
-            "prompt": prompt,
-            "use_lookup_cache": use_lookup_cache,
-            "ignore_unlabeled": ignore_unlabeled,
-            "partition_filter_mode": partition_filter_mode,
-        }
-        # Don't send partition_ids when partition_filter_mode is "ignore_partitions"
-        if partition_filter_mode != "ignore_partitions":
-            request_json["partition_ids"] = partition_id
-        response = client.POST(
-            endpoint,
-            params={"name_or_id": self.id},
-            json=request_json,
-            timeout=timeout_seconds,
-        )
 
-        if telemetry_on and any(p["prediction_id"] is None for p in response):
-            raise RuntimeError("Failed to save prediction to database.")
+        predictions: list[ClassificationPrediction] = []
+        for i in range(0, len(values), batch_size):
+            batch_values = values[i : i + batch_size]
+            batch_expected_labels = expected_labels[i : i + batch_size] if expected_labels else None
 
-        predictions = [
-            ClassificationPrediction(
-                prediction_id=prediction["prediction_id"],
-                label=prediction["label"],
-                label_name=prediction["label_name"],
-                score=None,
-                confidence=prediction["confidence"],
-                anomaly_score=prediction["anomaly_score"],
-                memoryset=self.memoryset,
-                model=self,
-                logits=prediction["logits"],
-                input_value=input_value,
+            request_json: ClassificationPredictionRequest = {
+                "input_values": batch_values,
+                "memoryset_override_name_or_id": self._memoryset_override_id,
+                "expected_labels": batch_expected_labels,
+                "tags": list(tags or set()),
+                "save_telemetry": telemetry_on,
+                "save_telemetry_synchronously": telemetry_sync,
+                "filters": cast(list[FilterItem], parsed_filters),
+                "prompt": prompt,
+                "use_lookup_cache": use_lookup_cache,
+                "ignore_unlabeled": ignore_unlabeled,
+                "partition_filter_mode": partition_filter_mode,
+            }
+            if partition_filter_mode != "ignore_partitions":
+                request_json["partition_ids"] = (
+                    partition_id[i : i + batch_size] if isinstance(partition_id, list) else partition_id
+                )
+
+            response = client.POST(
+                endpoint,
+                params={"name_or_id": self.id},
+                json=request_json,
+                timeout=timeout_seconds,
             )
-            for prediction, input_value in zip(response, value if isinstance(value, list) else [value])
-        ]
+
+            if telemetry_on and any(p["prediction_id"] is None for p in response):
+                raise RuntimeError("Failed to save some prediction to database.")
+
+            predictions.extend(
+                ClassificationPrediction(
+                    prediction_id=prediction["prediction_id"],
+                    label=prediction["label"],
+                    label_name=prediction["label_name"],
+                    score=None,
+                    confidence=prediction["confidence"],
+                    anomaly_score=prediction["anomaly_score"],
+                    memoryset=self.memoryset,
+                    model=self,
+                    logits=prediction["logits"],
+                    input_value=input_value,
+                )
+                for prediction, input_value in zip(response, batch_values)
+            )
+
         self._last_prediction_was_batch = isinstance(value, list)
         self._last_prediction = predictions[-1]
         return predictions if isinstance(value, list) else predictions[0]
@@ -542,6 +565,7 @@ class ClassificationModel:
         partition_filter_mode: Literal[
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
+        batch_size: int = 100,
     ) -> list[ClassificationPrediction]:
         pass
 
@@ -561,6 +585,7 @@ class ClassificationModel:
         partition_filter_mode: Literal[
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
+        batch_size: int = 100,
     ) -> ClassificationPrediction:
         pass
 
@@ -579,6 +604,7 @@ class ClassificationModel:
         partition_filter_mode: Literal[
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
+        batch_size: int = 100,
     ) -> list[ClassificationPrediction] | ClassificationPrediction:
         """
         Asynchronously predict label(s) for the given input value(s) grounded in similar memories
@@ -605,6 +631,8 @@ class ClassificationModel:
                 * `"include_global"`: Include global memories
                 * `"exclude_global"`: Exclude global memories
                 * `"only_global"`: Only include global memories
+            batch_size: Number of values to process in a single API call
+
         Returns:
             Label prediction or list of label predictions.
 
@@ -631,6 +659,8 @@ class ClassificationModel:
 
         if timeout_seconds <= 0:
             raise ValueError("timeout_seconds must be a positive integer")
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
 
         parsed_filters = [
             _parse_filter_item_from_tuple(filter) if isinstance(filter, tuple) else filter for filter in filters
@@ -639,10 +669,17 @@ class ClassificationModel:
         if any(_is_metric_column(filter[0]) for filter in filters):
             raise ValueError(f"Cannot filter on {filters} - telemetry filters are not supported for predictions")
 
+        # Convert to list for batching
+        values = value if isinstance(value, list) else [value]
+        if isinstance(expected_labels, list) and len(expected_labels) != len(values):
+            raise ValueError("Invalid input: \n\texpected_labels must be the same length as values")
+        if isinstance(partition_id, list) and len(partition_id) != len(values):
+            raise ValueError("Invalid input: \n\tpartition_id must be the same length as values")
+
         if isinstance(expected_labels, int):
-            expected_labels = [expected_labels]
+            expected_labels = [expected_labels] * len(values)
         elif isinstance(expected_labels, str):
-            expected_labels = [self.memoryset.label_names.index(expected_labels)]
+            expected_labels = [self.memoryset.label_names.index(expected_labels)] * len(values)
         elif isinstance(expected_labels, list):
             expected_labels = [
                 self.memoryset.label_names.index(label) if isinstance(label, str) else label
@@ -651,75 +688,89 @@ class ClassificationModel:
 
         telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
         client = OrcaAsyncClient._resolve_client()
-        request_json: ClassificationPredictionRequest = {
-            "input_values": value if isinstance(value, list) else [value],
-            "memoryset_override_name_or_id": self._memoryset_override_id,
-            "expected_labels": expected_labels,
-            "tags": list(tags or set()),
-            "save_telemetry": telemetry_on,
-            "save_telemetry_synchronously": telemetry_sync,
-            "filters": cast(list[FilterItem], parsed_filters),
-            "prompt": prompt,
-            "use_lookup_cache": use_lookup_cache,
-            "ignore_unlabeled": ignore_unlabeled,
-            "partition_filter_mode": partition_filter_mode,
-        }
-        # Don't send partition_ids when partition_filter_mode is "ignore_partitions"
-        if partition_filter_mode != "ignore_partitions":
-            request_json["partition_ids"] = partition_id
-        response = await client.POST(
-            "/gpu/classification_model/{name_or_id}/prediction",
-            params={"name_or_id": self.id},
-            json=request_json,
-            timeout=timeout_seconds,
-        )
 
-        if telemetry_on and any(p["prediction_id"] is None for p in response):
-            raise RuntimeError("Failed to save prediction to database.")
+        predictions: list[ClassificationPrediction] = []
+        for i in range(0, len(values), batch_size):
+            batch_values = values[i : i + batch_size]
+            batch_expected_labels = expected_labels[i : i + batch_size] if expected_labels else None
 
-        predictions = [
-            ClassificationPrediction(
-                prediction_id=prediction["prediction_id"],
-                label=prediction["label"],
-                label_name=prediction["label_name"],
-                score=None,
-                confidence=prediction["confidence"],
-                anomaly_score=prediction["anomaly_score"],
-                memoryset=self.memoryset,
-                model=self,
-                logits=prediction["logits"],
-                input_value=input_value,
+            request_json: ClassificationPredictionRequest = {
+                "input_values": batch_values,
+                "memoryset_override_name_or_id": self._memoryset_override_id,
+                "expected_labels": batch_expected_labels,
+                "tags": list(tags or set()),
+                "save_telemetry": telemetry_on,
+                "save_telemetry_synchronously": telemetry_sync,
+                "filters": cast(list[FilterItem], parsed_filters),
+                "prompt": prompt,
+                "use_lookup_cache": use_lookup_cache,
+                "ignore_unlabeled": ignore_unlabeled,
+                "partition_filter_mode": partition_filter_mode,
+            }
+            if partition_filter_mode != "ignore_partitions":
+                request_json["partition_ids"] = (
+                    partition_id[i : i + batch_size] if isinstance(partition_id, list) else partition_id
+                )
+            response = await client.POST(
+                "/gpu/classification_model/{name_or_id}/prediction",
+                params={"name_or_id": self.id},
+                json=request_json,
+                timeout=timeout_seconds,
             )
-            for prediction, input_value in zip(response, value if isinstance(value, list) else [value])
-        ]
+
+            if telemetry_on and any(p["prediction_id"] is None for p in response):
+                raise RuntimeError("Failed to save some prediction to database.")
+
+            predictions.extend(
+                ClassificationPrediction(
+                    prediction_id=prediction["prediction_id"],
+                    label=prediction["label"],
+                    label_name=prediction["label_name"],
+                    score=None,
+                    confidence=prediction["confidence"],
+                    anomaly_score=prediction["anomaly_score"],
+                    memoryset=self.memoryset,
+                    model=self,
+                    logits=prediction["logits"],
+                    input_value=input_value,
+                )
+                for prediction, input_value in zip(response, batch_values)
+            )
+
         self._last_prediction_was_batch = isinstance(value, list)
         self._last_prediction = predictions[-1]
         return predictions if isinstance(value, list) else predictions[0]
 
     def predictions(
         self,
-        limit: int = 100,
+        limit: int | None = None,
         offset: int = 0,
         tag: str | None = None,
         sort: list[tuple[Literal["anomaly_score", "confidence", "timestamp"], Literal["asc", "desc"]]] = [],
         expected_label_match: bool | None = None,
+        batch_size: int = 100,
     ) -> list[ClassificationPrediction]:
         """
         Get a list of predictions made by this model
 
         Params:
-            limit: Optional maximum number of predictions to return
+            limit: Maximum number of predictions to return. If `None`, returns all predictions
+                by automatically paginating through results.
             offset: Optional offset of the first prediction to return
             tag: Optional tag to filter predictions by
             sort: Optional list of columns and directions to sort the predictions by.
                 Predictions can be sorted by `timestamp` or `confidence`.
             expected_label_match: Optional filter to only include predictions where the expected
                 label does (`True`) or doesn't (`False`) match the predicted label
+            batch_size: Number of predictions to fetch in a single API call
 
         Returns:
             List of label predictions
 
         Examples:
+            Get all predictions with a specific tag:
+            >>> predictions = model.predictions(tag="evaluation")
+
             Get the last 3 predictions:
             >>> predictions = model.predictions(limit=3, sort=[("timestamp", "desc")])
             [
@@ -737,33 +788,61 @@ class ClassificationModel:
             >>> predictions = model.predictions(expected_label_match=False)
             [ClassificationPrediction({label: <positive: 1>, confidence: 0.95, anomaly_score: 0.1, input_value: 'I am happy', expected_label: 0})]
         """
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
+        if limit == 0:
+            return []
+
         client = OrcaClient._resolve_client()
-        predictions = client.POST(
-            "/telemetry/prediction",
-            json={
+        all_predictions: list[ClassificationPrediction] = []
+
+        if limit is not None and limit < batch_size:
+            pages = [(offset, limit)]
+        else:
+            # automatically paginate the requests if necessary
+            total = client.POST(
+                "/telemetry/prediction/count",
+                json={
+                    "model_id": self.id,
+                    "tag": tag,
+                    "expected_label_match": expected_label_match,
+                },
+            )
+            max_limit = max(total - offset, 0)
+            limit = min(limit, max_limit) if limit is not None else max_limit
+            pages = [(o, min(batch_size, limit - (o - offset))) for o in range(offset, offset + limit, batch_size)]
+
+        for current_offset, current_limit in pages:
+            request_json: ListPredictionsRequest = {
                 "model_id": self.id,
-                "limit": limit,
-                "offset": offset,
-                "sort": [list(sort_item) for sort_item in sort],
+                "limit": current_limit,
+                "offset": current_offset,
                 "tag": tag,
                 "expected_label_match": expected_label_match,
-            },
-        )
-        return [
-            ClassificationPrediction(
-                prediction_id=prediction["prediction_id"],
-                label=prediction["label"],
-                label_name=prediction["label_name"],
-                score=None,
-                confidence=prediction["confidence"],
-                anomaly_score=prediction["anomaly_score"],
-                memoryset=self.memoryset,
-                model=self,
-                telemetry=prediction,
+            }
+            if sort:
+                request_json["sort"] = sort
+            response = client.POST(
+                "/telemetry/prediction",
+                json=request_json,
             )
-            for prediction in predictions
-            if "label" in prediction
-        ]
+            all_predictions.extend(
+                ClassificationPrediction(
+                    prediction_id=prediction["prediction_id"],
+                    label=prediction["label"],
+                    label_name=prediction["label_name"],
+                    score=None,
+                    confidence=prediction["confidence"],
+                    anomaly_score=prediction["anomaly_score"],
+                    memoryset=self.memoryset,
+                    model=self,
+                    telemetry=prediction,
+                )
+                for prediction in response
+                if "label" in prediction
+            )
+
+        return all_predictions
 
     def _evaluate_datasource(
         self,
@@ -861,6 +940,7 @@ class ClassificationModel:
             logits=[p.logits for p in predictions],
             anomaly_scores=[p.anomaly_score for p in predictions],
             include_curves=True,
+            include_confusion_matrix=True,
         )
 
     @overload

orca_sdk/classification_model_test.py

@@ -218,13 +218,17 @@ def test_evaluate_dataset_with_nones_raises_error(classification_model: Classifi
 
 
 def test_evaluate_with_telemetry(classification_model: ClassificationModel, eval_dataset: Dataset):
-    result = classification_model.evaluate(eval_dataset, record_predictions=True, tags={"test"})
+    result = classification_model.evaluate(eval_dataset, record_predictions=True, tags={"test"}, batch_size=2)
     assert result is not None
     assert isinstance(result, ClassificationMetrics)
-    predictions = classification_model.predictions(tag="test")
+    predictions = classification_model.predictions(tag="test", batch_size=100, sort=[("timestamp", "asc")])
     assert len(predictions) == 4
     assert all(p.tags == {"test"} for p in predictions)
-    assert all(p.expected_label == l for p, l in zip(predictions, eval_dataset["label"]))
+    prediction_expected_labels = [p.expected_label if p.expected_label is not None else -1 for p in predictions]
+    eval_expected_labels = list(eval_dataset["label"])
+    assert all(
+        p == l for p, l in zip(prediction_expected_labels, eval_expected_labels)
+    ), f"Prediction expected labels: {prediction_expected_labels} do not match eval expected labels: {eval_expected_labels}"
 
 
 def test_evaluate_with_partition_column_dataset(partitioned_classification_model: ClassificationModel):
@@ -361,7 +365,7 @@ def test_evaluate_with_partition_column_datasource(partitioned_classification_mo
 
 
 def test_predict(classification_model: ClassificationModel, label_names: list[str]):
-    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"])
+    predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"], batch_size=1)
     assert len(predictions) == 2
     assert predictions[0].prediction_id is not None
     assert predictions[1].prediction_id is not None