orca-sdk 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

orca_sdk/classification_model.py

@@ -12,11 +12,10 @@ from ._utils.common import UNSET, CreateMode, DropMode
 from .async_client import OrcaAsyncClient
 from .client import (
     BootstrapClassificationModelMeta,
-    BootstrapClassificationModelResult,
-    ClassificationEvaluationRequest,
+    BootstrapLabeledMemoryDataResult,
     ClassificationModelMetadata,
+    ClassificationPredictionRequest,
     OrcaClient,
-    PostClassificationModelByModelNameOrIdEvaluationParams,
     PredictiveModelUpdate,
     RACHeadType,
 )
@@ -42,7 +41,7 @@ class BootstrappedClassificationModel:
     datasource: Datasource | None
     memoryset: LabeledMemoryset | None
     classification_model: ClassificationModel | None
-    agent_output: BootstrapClassificationModelResult | None
+    agent_output: BootstrapLabeledMemoryDataResult | None
 
     def __init__(self, metadata: BootstrapClassificationModelMeta):
         self.datasource = Datasource.open(metadata["datasource_meta"]["id"])
@@ -358,6 +357,10 @@ class ClassificationModel:
         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
@@ -374,6 +377,10 @@ class ClassificationModel:
         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
@@ -389,6 +396,10 @@ class ClassificationModel:
         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:
         """
@@ -410,6 +421,12 @@ class ClassificationModel:
             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:
@@ -463,21 +480,26 @@ 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={
-                "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,
-            },
+            json=request_json,
             timeout=timeout_seconds,
         )
 
@@ -515,6 +537,10 @@ class ClassificationModel:
         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
 
@@ -530,6 +556,10 @@ class ClassificationModel:
         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
 
@@ -544,6 +574,10 @@ class ClassificationModel:
         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
@@ -564,7 +598,12 @@ class ClassificationModel:
             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.
 
@@ -611,21 +650,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,
-                "ignore_unlabeled": ignore_unlabeled,
-            },
+            json=request_json,
             timeout=timeout_seconds,
         )
 
@@ -730,6 +774,10 @@ class ClassificationModel:
         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(
@@ -744,6 +792,8 @@ class ClassificationModel:
                 "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,
             },
         )
 
@@ -780,6 +830,10 @@ class ClassificationModel:
         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")
@@ -796,6 +850,8 @@ class ClassificationModel:
                 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,
             )
         ]
 
@@ -813,12 +869,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
 
@@ -829,12 +889,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
 
@@ -844,12 +908,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
@@ -858,13 +926,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
 
@@ -888,6 +961,8 @@ class ClassificationModel:
                 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(
@@ -898,6 +973,8 @@ class ClassificationModel:
                 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)}")

orca_sdk/classification_model_test.py

@@ -187,18 +187,24 @@ def test_evaluate(classification_model, eval_datasource: Datasource, eval_datase
     assert -1.0 <= result.anomaly_score_variance <= 1.0
 
     assert result.pr_auc is not None
-    assert np.allclose(result.pr_auc, 0.75)
+    assert np.allclose(result.pr_auc, 0.83333)
     assert result.pr_curve is not None
-    assert np.allclose(result.pr_curve["thresholds"], [0.0, 0.0, 0.8155114054679871, 0.834095299243927])
-    assert np.allclose(result.pr_curve["precisions"], [0.5, 0.5, 1.0, 1.0])
-    assert np.allclose(result.pr_curve["recalls"], [1.0, 0.5, 0.5, 0.0])
+    assert np.allclose(
+        result.pr_curve["thresholds"],
+        [0.0, 0.3021204173564911, 0.30852025747299194, 0.6932827234268188, 0.6972201466560364],
+    )
+    assert np.allclose(result.pr_curve["precisions"], [0.5, 0.666666, 0.5, 1.0, 1.0])
+    assert np.allclose(result.pr_curve["recalls"], [1.0, 1.0, 0.5, 0.5, 0.0])
 
     assert result.roc_auc is not None
-    assert np.allclose(result.roc_auc, 0.625)
+    assert np.allclose(result.roc_auc, 0.75)
     assert result.roc_curve is not None
-    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["thresholds"],
+        [0.3021204173564911, 0.30852025747299194, 0.6932827234268188, 0.6972201466560364, 1.0],
+    )
+    assert np.allclose(result.roc_curve["false_positive_rates"], [1.0, 0.5, 0.5, 0.0, 0.0])
+    assert np.allclose(result.roc_curve["true_positive_rates"], [1.0, 1.0, 0.5, 0.5, 0.0])
 
 
 def test_evaluate_datasource_with_nones_raises_error(classification_model: ClassificationModel, datasource: Datasource):
@@ -221,6 +227,139 @@ def test_evaluate_with_telemetry(classification_model: ClassificationModel, eval
     assert all(p.expected_label == l for p, l in zip(predictions, eval_dataset["label"]))
 
 
+def test_evaluate_with_partition_column_dataset(partitioned_classification_model: ClassificationModel):
+    """Test evaluate with partition_column on a Dataset"""
+    # Create a test dataset with partition_id column
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "label": 0, "partition_id": "p1"},
+            {"value": "cats are cute", "label": 1, "partition_id": "p1"},
+            {"value": "homemade soup recipes", "label": 0, "partition_id": "p2"},
+            {"value": "cats purr when happy", "label": 1, "partition_id": "p2"},
+        ]
+    )
+
+    # Evaluate with partition_column
+    result = partitioned_classification_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="exclude_global",
+    )
+    assert result is not None
+    assert isinstance(result, ClassificationMetrics)
+    assert isinstance(result.accuracy, float)
+    assert isinstance(result.f1_score, float)
+    assert isinstance(result.loss, float)
+
+
+def test_evaluate_with_partition_column_include_global(partitioned_classification_model: ClassificationModel):
+    """Test evaluate with partition_column and include_global mode"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "label": 0, "partition_id": "p1"},
+            {"value": "cats are cute", "label": 1, "partition_id": "p1"},
+        ]
+    )
+
+    # Evaluate with partition_column and include_global (default)
+    result = partitioned_classification_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="include_global",
+    )
+    assert result is not None
+    assert isinstance(result, ClassificationMetrics)
+
+
+def test_evaluate_with_partition_column_exclude_global(partitioned_classification_model: ClassificationModel):
+    """Test evaluate with partition_column and exclude_global mode"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "label": 0, "partition_id": "p1"},
+            {"value": "cats are cute", "label": 1, "partition_id": "p1"},
+        ]
+    )
+
+    # Evaluate with partition_column and exclude_global
+    result = partitioned_classification_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="exclude_global",
+    )
+    assert result is not None
+    assert isinstance(result, ClassificationMetrics)
+
+
+def test_evaluate_with_partition_column_only_global(partitioned_classification_model: ClassificationModel):
+    """Test evaluate with partition_filter_mode only_global"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "cats are independent animals", "label": 1, "partition_id": None},
+            {"value": "i love the beach", "label": 1, "partition_id": None},
+        ]
+    )
+
+    # Evaluate with only_global mode
+    result = partitioned_classification_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="only_global",
+    )
+    assert result is not None
+    assert isinstance(result, ClassificationMetrics)
+
+
+def test_evaluate_with_partition_column_ignore_partitions(partitioned_classification_model: ClassificationModel):
+    """Test evaluate with partition_filter_mode ignore_partitions"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "label": 0, "partition_id": "p1"},
+            {"value": "cats are cute", "label": 1, "partition_id": "p2"},
+        ]
+    )
+
+    # Evaluate with ignore_partitions mode
+    result = partitioned_classification_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="ignore_partitions",
+    )
+    assert result is not None
+    assert isinstance(result, ClassificationMetrics)
+
+
+@pytest.mark.parametrize("data_type", ["dataset", "datasource"])
+def test_evaluate_with_partition_column_datasource(partitioned_classification_model: ClassificationModel, data_type):
+    """Test evaluate with partition_column on a Datasource"""
+    # Create a test datasource with partition_id column
+    eval_data_with_partition = [
+        {"value": "soup is good", "label": 0, "partition_id": "p1"},
+        {"value": "cats are cute", "label": 1, "partition_id": "p1"},
+        {"value": "homemade soup recipes", "label": 0, "partition_id": "p2"},
+        {"value": "cats purr when happy", "label": 1, "partition_id": "p2"},
+    ]
+
+    if data_type == "dataset":
+        eval_data = Dataset.from_list(eval_data_with_partition)
+        result = partitioned_classification_model.evaluate(
+            eval_data,
+            partition_column="partition_id",
+            partition_filter_mode="exclude_global",
+        )
+    else:
+        eval_datasource = Datasource.from_list("eval_datasource_with_partition", eval_data_with_partition)
+        result = partitioned_classification_model.evaluate(
+            eval_datasource,
+            partition_column="partition_id",
+            partition_filter_mode="exclude_global",
+        )
+
+    assert result is not None
+    assert isinstance(result, ClassificationMetrics)
+    assert isinstance(result.accuracy, float)
+    assert isinstance(result.f1_score, float)
+
+
 def test_predict(classification_model: ClassificationModel, label_names: list[str]):
     predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"])
     assert len(predictions) == 2
@@ -284,6 +423,186 @@ def test_predict_constraint_violation(readonly_memoryset: LabeledMemoryset):
         model.predict("test")
 
 
+def test_predict_with_partition_id(partitioned_classification_model: ClassificationModel, label_names: list[str]):
+    """Test predict with a specific partition_id"""
+    # Predict with partition_id p1 - should use memories from p1
+    prediction = partitioned_classification_model.predict(
+        "soup", partition_id="p1", partition_filter_mode="exclude_global"
+    )
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+    assert prediction.logits is not None
+    assert len(prediction.logits) == 2
+
+    # Predict with partition_id p2 - should use memories from p2
+    prediction_p2 = partitioned_classification_model.predict(
+        "cats", partition_id="p2", partition_filter_mode="exclude_global"
+    )
+    assert prediction_p2.label is not None
+    assert prediction_p2.label_name in label_names
+    assert 0 <= prediction_p2.confidence <= 1
+
+
+def test_predict_with_partition_id_include_global(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test predict with partition_id and include_global mode (default)"""
+    # Predict with partition_id p1 and include_global (default) - should include both p1 and global memories
+    prediction = partitioned_classification_model.predict(
+        "soup", partition_id="p1", partition_filter_mode="include_global"
+    )
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_with_partition_id_exclude_global(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test predict with partition_id and exclude_global mode"""
+    # Predict with partition_id p1 and exclude_global - should only use p1 memories
+    prediction = partitioned_classification_model.predict(
+        "soup", partition_id="p1", partition_filter_mode="exclude_global"
+    )
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_with_partition_id_only_global(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test predict with partition_filter_mode only_global"""
+    # Predict with only_global mode - should only use global memories
+    prediction = partitioned_classification_model.predict("cats", partition_filter_mode="only_global")
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_with_partition_id_ignore_partitions(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test predict with partition_filter_mode ignore_partitions"""
+    # Predict with ignore_partitions mode - should ignore partition filtering
+    prediction = partitioned_classification_model.predict("soup", partition_filter_mode="ignore_partitions")
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_batch_with_partition_id(partitioned_classification_model: ClassificationModel, label_names: list[str]):
+    """Test batch predict with partition_id"""
+    # Batch predict with partition_id p1
+    predictions = partitioned_classification_model.predict(
+        ["soup is good", "cats are cute"],
+        partition_id="p1",
+        partition_filter_mode="exclude_global",
+    )
+    assert len(predictions) == 2
+    assert all(p.label is not None for p in predictions)
+    assert all(p.label_name in label_names for p in predictions)
+    assert all(0 <= p.confidence <= 1 for p in predictions)
+    assert all(p.logits is not None and len(p.logits) == 2 for p in predictions)
+
+
+def test_predict_with_partition_id_and_filters(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test predict with partition_id and filters"""
+    # Predict with partition_id and filters
+    prediction = partitioned_classification_model.predict(
+        "soup",
+        partition_id="p1",
+        partition_filter_mode="exclude_global",
+        filters=[("key", "==", "g1")],
+    )
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_batch_with_list_of_partition_ids(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test batch predict with a list of partition_ids (one for each query input)"""
+    # Batch predict with a list of partition_ids - one for each input
+    # First input uses p1, second input uses p2
+    predictions = partitioned_classification_model.predict(
+        ["soup is good", "cats are cute"],
+        partition_id=["p1", "p2"],
+        partition_filter_mode="exclude_global",
+    )
+    assert len(predictions) == 2
+    assert all(p.label is not None for p in predictions)
+    assert all(p.label_name in label_names for p in predictions)
+    assert all(0 <= p.confidence <= 1 for p in predictions)
+    assert all(p.logits is not None and len(p.logits) == 2 for p in predictions)
+
+    # Verify that predictions were made using the correct partitions
+    # Each prediction should use memories from its respective partition
+    assert predictions[0].input_value == "soup is good"
+    assert predictions[1].input_value == "cats are cute"
+
+
+@pytest.mark.asyncio
+async def test_predict_async_with_partition_id(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test async predict with partition_id"""
+    # Async predict with partition_id p1
+    prediction = await partitioned_classification_model.apredict(
+        "soup", partition_id="p1", partition_filter_mode="exclude_global"
+    )
+    assert prediction.label is not None
+    assert prediction.label_name in label_names
+    assert 0 <= prediction.confidence <= 1
+    assert prediction.logits is not None
+    assert len(prediction.logits) == 2
+
+
+@pytest.mark.asyncio
+async def test_predict_async_batch_with_partition_id(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test async batch predict with partition_id"""
+    # Async batch predict with partition_id p1
+    predictions = await partitioned_classification_model.apredict(
+        ["soup is good", "cats are cute"],
+        partition_id="p1",
+        partition_filter_mode="exclude_global",
+    )
+    assert len(predictions) == 2
+    assert all(p.label is not None for p in predictions)
+    assert all(p.label_name in label_names for p in predictions)
+    assert all(0 <= p.confidence <= 1 for p in predictions)
+
+
+@pytest.mark.asyncio
+async def test_predict_async_batch_with_list_of_partition_ids(
+    partitioned_classification_model: ClassificationModel, label_names: list[str]
+):
+    """Test async batch predict with a list of partition_ids (one for each query input)"""
+    # Async batch predict with a list of partition_ids - one for each input
+    # First input uses p1, second input uses p2
+    predictions = await partitioned_classification_model.apredict(
+        ["soup is good", "cats are cute"],
+        partition_id=["p1", "p2"],
+        partition_filter_mode="exclude_global",
+    )
+    assert len(predictions) == 2
+    assert all(p.label is not None for p in predictions)
+    assert all(p.label_name in label_names for p in predictions)
+    assert all(0 <= p.confidence <= 1 for p in predictions)
+    assert all(p.logits is not None and len(p.logits) == 2 for p in predictions)
+
+    # Verify that predictions were made using the correct partitions
+    # Each prediction should use memories from its respective partition
+    assert predictions[0].input_value == "soup is good"
+    assert predictions[1].input_value == "cats are cute"
+
+
 def test_record_prediction_feedback(classification_model: ClassificationModel):
     predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"])
     expected_labels = [0, 1]