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

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]