orca-sdk 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

orca_sdk/classification_model.py

@@ -3,14 +3,7 @@ from __future__ import annotations
 import logging
 from contextlib import contextmanager
 from datetime import datetime
-from typing import (
-    Any,
-    Generator,
-    Iterable,
-    Literal,
-    cast,
-    overload,
-)
+from typing import Any, Generator, Iterable, Literal, cast, overload
 
 from datasets import Dataset
 
@@ -20,8 +13,11 @@ from .async_client import OrcaAsyncClient
 from .client import (
     BootstrapClassificationModelMeta,
     BootstrapClassificationModelResult,
+    ClassificationEvaluationRequest,
     ClassificationModelMetadata,
+    ClassificationPredictionRequest,
     OrcaClient,
+    PostClassificationModelByModelNameOrIdEvaluationParams,
     PredictiveModelUpdate,
     RACHeadType,
 )
@@ -207,7 +203,12 @@ class ClassificationModel:
                 raise ValueError(f"Model with name {name} already exists")
             elif if_exists == "open":
                 existing = cls.open(name)
-                for attribute in {"head_type", "memory_lookup_count", "num_classes", "min_memory_weight"}:
+                for attribute in {
+                    "head_type",
+                    "memory_lookup_count",
+                    "num_classes",
+                    "min_memory_weight",
+                }:
                     local_attribute = locals()[attribute]
                     existing_attribute = getattr(existing, attribute)
                     if local_attribute is not None and local_attribute != existing_attribute:
@@ -357,6 +358,12 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        partition_id: str | list[str | None] | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
+        use_gpu: bool = True,
     ) -> list[ClassificationPrediction]:
         pass
 
@@ -371,6 +378,12 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        partition_id: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
+        use_gpu: bool = True,
     ) -> ClassificationPrediction:
         pass
 
@@ -384,6 +397,12 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        partition_id: str | None | list[str | None] = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
+        use_gpu: bool = True,
     ) -> list[ClassificationPrediction] | ClassificationPrediction:
         """
         Predict label(s) for the given input value(s) grounded in similar memories
@@ -402,6 +421,15 @@ class ClassificationModel:
             prompt: Optional prompt to use for instruction-tuned embedding models
             use_lookup_cache: Whether to use cached lookup results for faster predictions
             timeout_seconds: Timeout in seconds for the request, defaults to 10 seconds
+            ignore_unlabeled: If True, only use labeled memories during lookup.
+                If False (default), allow unlabeled memories when necessary.
+            partition_id: Optional partition ID(s) to use during memory lookup
+            partition_filter_mode: Optional partition filter mode to use for the prediction(s). One of
+                * `"ignore_partitions"`: Ignore partitions
+                * `"include_global"`: Include global memories
+                * `"exclude_global"`: Exclude global memories
+                * `"only_global"`: Only include global memories
+            use_gpu: Whether to use GPU for the prediction (defaults to True)
 
         Returns:
             Label prediction or list of label predictions
@@ -447,22 +475,33 @@ class ClassificationModel:
                 for label in expected_labels
             ]
 
+        if use_gpu:
+            endpoint = "/gpu/classification_model/{name_or_id}/prediction"
+        else:
+            endpoint = "/classification_model/{name_or_id}/prediction"
+
         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(
-            "/gpu/classification_model/{name_or_id}/prediction",
+            endpoint,
             params={"name_or_id": self.id},
-            json={
-                "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,
-            },
+            json=request_json,
             timeout=timeout_seconds,
         )
 
@@ -499,6 +538,11 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        partition_id: str | list[str | None] | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> list[ClassificationPrediction]:
         pass
 
@@ -513,6 +557,11 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        partition_id: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> ClassificationPrediction:
         pass
 
@@ -526,6 +575,11 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        partition_id: str | None | list[str | None] = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> list[ClassificationPrediction] | ClassificationPrediction:
         """
         Asynchronously predict label(s) for the given input value(s) grounded in similar memories
@@ -544,7 +598,14 @@ class ClassificationModel:
             prompt: Optional prompt to use for instruction-tuned embedding models
             use_lookup_cache: Whether to use cached lookup results for faster predictions
             timeout_seconds: Timeout in seconds for the request, defaults to 10 seconds
-
+            ignore_unlabeled: If True, only use labeled memories during lookup.
+                If False (default), allow unlabeled memories when necessary.
+            partition_id: Optional partition ID(s) to use during memory lookup
+            partition_filter_mode: Optional partition filter mode to use for the prediction(s). One of
+                * `"ignore_partitions"`: Ignore partitions
+                * `"include_global"`: Include global memories
+                * `"exclude_global"`: Exclude global memories
+                * `"only_global"`: Only include global memories
         Returns:
             Label prediction or list of label predictions.
 
@@ -591,20 +652,26 @@ 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={
-                "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,
-            },
+            json=request_json,
             timeout=timeout_seconds,
         )
 
@@ -706,7 +773,13 @@ class ClassificationModel:
         label_column: str,
         record_predictions: bool,
         tags: set[str] | None,
+        subsample: int | float | None,
         background: bool = False,
+        ignore_unlabeled: bool = False,
+        partition_column: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> ClassificationMetrics | Job[ClassificationMetrics]:
         client = OrcaClient._resolve_client()
         response = client.POST(
@@ -719,14 +792,18 @@ class ClassificationModel:
                 "memoryset_override_name_or_id": self._memoryset_override_id,
                 "record_telemetry": record_predictions,
                 "telemetry_tags": list(tags) if tags else None,
+                "subsample": subsample,
+                "ignore_unlabeled": ignore_unlabeled,
+                "datasource_partition_column": partition_column,
+                "partition_filter_mode": partition_filter_mode,
             },
         )
 
         def get_value():
             client = OrcaClient._resolve_client()
             res = client.GET(
-                "/classification_model/{model_name_or_id}/evaluation/{task_id}",
-                params={"model_name_or_id": self.id, "task_id": response["task_id"]},
+                "/classification_model/{model_name_or_id}/evaluation/{job_id}",
+                params={"model_name_or_id": self.id, "job_id": response["job_id"]},
             )
             assert res["result"] is not None
             return ClassificationMetrics(
@@ -743,7 +820,7 @@ class ClassificationModel:
                 roc_curve=res["result"].get("roc_curve"),
             )
 
-        job = Job(response["task_id"], get_value)
+        job = Job(response["job_id"], get_value)
         return job if background else job.result()
 
     def _evaluate_dataset(
@@ -754,6 +831,11 @@ class ClassificationModel:
         record_predictions: bool,
         tags: set[str],
         batch_size: int,
+        ignore_unlabeled: bool,
+        partition_column: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> ClassificationMetrics:
         if len(dataset) == 0:
             raise ValueError("Evaluation dataset cannot be empty")
@@ -769,6 +851,9 @@ class ClassificationModel:
                 expected_labels=dataset[i : i + batch_size][label_column],
                 tags=tags,
                 save_telemetry="sync" if record_predictions else "off",
+                ignore_unlabeled=ignore_unlabeled,
+                partition_id=dataset[i : i + batch_size][partition_column] if partition_column else None,
+                partition_filter_mode=partition_filter_mode,
             )
         ]
 
@@ -786,10 +871,16 @@ class ClassificationModel:
         *,
         value_column: str = "value",
         label_column: str = "label",
+        partition_column: str | None = None,
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        subsample: int | float | None = None,
         background: Literal[True],
+        ignore_unlabeled: bool = False,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> Job[ClassificationMetrics]:
         pass
 
@@ -800,10 +891,16 @@ class ClassificationModel:
         *,
         value_column: str = "value",
         label_column: str = "label",
+        partition_column: str | None = None,
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        subsample: int | float | None = None,
         background: Literal[False] = False,
+        ignore_unlabeled: bool = False,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> ClassificationMetrics:
         pass
 
@@ -813,10 +910,16 @@ class ClassificationModel:
         *,
         value_column: str = "value",
         label_column: str = "label",
+        partition_column: str | None = None,
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        subsample: int | float | None = None,
         background: bool = False,
+        ignore_unlabeled: bool = False,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> ClassificationMetrics | Job[ClassificationMetrics]:
         """
         Evaluate the classification model on a given dataset or datasource
@@ -825,11 +928,18 @@ class ClassificationModel:
             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
+            partition_column: Optional name of the column that contains the partition IDs
             record_predictions: Whether to record [`ClassificationPrediction`][orca_sdk.telemetry.ClassificationPrediction]s for analysis
             tags: Optional tags to add to the recorded [`ClassificationPrediction`][orca_sdk.telemetry.ClassificationPrediction]s
             batch_size: Batch size for processing Dataset inputs (only used when input is a Dataset)
+            subsample: Optional number (int) of rows to sample or fraction (float in (0, 1]) of data to sample for evaluation.
             background: Whether to run the operation in the background and return a job handle
-
+            ignore_unlabeled: If True, only use labeled memories during lookup. If False (default), allow unlabeled memories
+            partition_filter_mode: Optional partition filter mode to use for the evaluation. One of
+                * `"ignore_partitions"`: Ignore partitions
+                * `"include_global"`: Include global memories
+                * `"exclude_global"`: Exclude global memories
+                * `"only_global"`: Only include global memories
         Returns:
             EvaluationResult containing metrics including accuracy, F1 score, ROC AUC, PR AUC, and anomaly score statistics
 
@@ -850,7 +960,11 @@ class ClassificationModel:
                 label_column=label_column,
                 record_predictions=record_predictions,
                 tags=tags,
+                subsample=subsample,
                 background=background,
+                ignore_unlabeled=ignore_unlabeled,
+                partition_column=partition_column,
+                partition_filter_mode=partition_filter_mode,
             )
         elif isinstance(data, Dataset):
             return self._evaluate_dataset(
@@ -860,6 +974,9 @@ class ClassificationModel:
                 record_predictions=record_predictions,
                 tags=tags,
                 batch_size=batch_size,
+                ignore_unlabeled=ignore_unlabeled,
+                partition_column=partition_column,
+                partition_filter_mode=partition_filter_mode,
             )
         else:
             raise ValueError(f"Invalid data type: {type(data)}")
@@ -961,11 +1078,9 @@ class ClassificationModel:
 
         def get_result() -> BootstrappedClassificationModel:
             client = OrcaClient._resolve_client()
-            res = client.GET(
-                "/agents/bootstrap_classification_model/{task_id}", params={"task_id": response["task_id"]}
-            )
+            res = client.GET("/agents/bootstrap_classification_model/{job_id}", params={"job_id": response["job_id"]})
             assert res["result"] is not None
             return BootstrappedClassificationModel(res["result"])
 
-        job = Job(response["task_id"], get_result)
+        job = Job(response["job_id"], get_result)
         return job if background else job.result()