orca-sdk 0.1.0__py3-none-any.whl

orca_sdk/classification_model_test.py

@@ -0,0 +1,266 @@
+from uuid import uuid4
+
+import pytest
+from datasets.arrow_dataset import Dataset
+
+from .classification_model import ClassificationModel
+from .datasource import Datasource
+from .embedding_model import PretrainedEmbeddingModel
+from .labeled_memoryset import LabeledMemoryset
+
+
+def test_create_model(model: ClassificationModel, memoryset: LabeledMemoryset):
+    assert model is not None
+    assert model.name == "test_model"
+    assert model.memoryset == memoryset
+    assert model.num_classes == 2
+    assert model.memory_lookup_count == 3
+
+
+def test_create_model_already_exists_error(memoryset, model: ClassificationModel):
+    with pytest.raises(ValueError):
+        ClassificationModel.create("test_model", memoryset)
+    with pytest.raises(ValueError):
+        ClassificationModel.create("test_model", memoryset, if_exists="error")
+
+
+def test_create_model_already_exists_return(memoryset, model: ClassificationModel):
+    with pytest.raises(ValueError):
+        ClassificationModel.create("test_model", memoryset, if_exists="open", head_type="MMOE")
+
+    with pytest.raises(ValueError):
+        ClassificationModel.create("test_model", memoryset, if_exists="open", memory_lookup_count=37)
+
+    with pytest.raises(ValueError):
+        ClassificationModel.create("test_model", memoryset, if_exists="open", num_classes=19)
+
+    with pytest.raises(ValueError):
+        ClassificationModel.create("test_model", memoryset, if_exists="open", min_memory_weight=0.77)
+
+    new_model = ClassificationModel.create("test_model", memoryset, if_exists="open")
+    assert new_model is not None
+    assert new_model.name == "test_model"
+    assert new_model.memoryset == memoryset
+    assert new_model.num_classes == 2
+    assert new_model.memory_lookup_count == 3
+
+
+def test_create_model_unauthenticated(unauthenticated, memoryset: LabeledMemoryset):
+    with pytest.raises(ValueError, match="Invalid API key"):
+        ClassificationModel.create("test_model", memoryset)
+
+
+def test_get_model(model: ClassificationModel):
+    fetched_model = ClassificationModel.open(model.name)
+    assert fetched_model is not None
+    assert fetched_model.id == model.id
+    assert fetched_model.name == model.name
+    assert fetched_model.num_classes == 2
+    assert fetched_model.memory_lookup_count == 3
+    assert fetched_model == model
+
+
+def test_get_model_unauthenticated(unauthenticated):
+    with pytest.raises(ValueError, match="Invalid API key"):
+        ClassificationModel.open("test_model")
+
+
+def test_get_model_invalid_input():
+    with pytest.raises(ValueError, match="Invalid input"):
+        ClassificationModel.open("not valid id")
+
+
+def test_get_model_not_found():
+    with pytest.raises(LookupError):
+        ClassificationModel.open(str(uuid4()))
+
+
+def test_get_model_unauthorized(unauthorized, model: ClassificationModel):
+    with pytest.raises(LookupError):
+        ClassificationModel.open(model.name)
+
+
+def test_list_models(model: ClassificationModel):
+    models = ClassificationModel.all()
+    assert len(models) > 0
+    assert any(model.name == model.name for model in models)
+
+
+def test_list_models_unauthenticated(unauthenticated):
+    with pytest.raises(ValueError, match="Invalid API key"):
+        ClassificationModel.all()
+
+
+def test_list_models_unauthorized(unauthorized, model: ClassificationModel):
+    assert ClassificationModel.all() == []
+
+
+def test_delete_model(memoryset: LabeledMemoryset):
+    ClassificationModel.create("model_to_delete", LabeledMemoryset.open(memoryset.name))
+    assert ClassificationModel.open("model_to_delete")
+    ClassificationModel.drop("model_to_delete")
+    with pytest.raises(LookupError):
+        ClassificationModel.open("model_to_delete")
+
+
+def test_delete_model_unauthenticated(unauthenticated, model: ClassificationModel):
+    with pytest.raises(ValueError, match="Invalid API key"):
+        ClassificationModel.drop(model.name)
+
+
+def test_delete_model_not_found():
+    with pytest.raises(LookupError):
+        ClassificationModel.drop(str(uuid4()))
+    # ignores error if specified
+    ClassificationModel.drop(str(uuid4()), if_not_exists="ignore")
+
+
+def test_delete_model_unauthorized(unauthorized, model: ClassificationModel):
+    with pytest.raises(LookupError):
+        ClassificationModel.drop(model.name)
+
+
+@pytest.mark.flaky
+def test_delete_memoryset_before_model_constraint_violation(hf_dataset):
+    memoryset = LabeledMemoryset.from_hf_dataset("test_memoryset_delete_before_model", hf_dataset, value_column="text")
+    ClassificationModel.create("test_model_delete_before_memoryset", memoryset)
+    with pytest.raises(RuntimeError):
+        LabeledMemoryset.drop(memoryset.id)
+
+
+def test_evaluate(model):
+    eval_datasource = Datasource.from_list(
+        "eval_datasource",
+        [
+            {"text": "chicken noodle soup is the best", "label": 1},
+            {"text": "cats are cute", "label": 0},
+            {"text": "soup is great for the winter", "label": 0},
+            {"text": "i love cats", "label": 1},
+        ],
+    )
+    result = model.evaluate(eval_datasource, value_column="text")
+    assert result is not None
+    assert isinstance(result["accuracy"], float)
+    assert isinstance(result["f1_score"], float)
+    assert isinstance(result["loss"], float)
+
+
+def test_evaluate_with_telemetry(model):
+    samples = [
+        {"text": "chicken noodle soup is the best", "label": 1},
+        {"text": "cats are cute", "label": 0},
+    ]
+    eval_datasource = Datasource.from_list("eval_datasource_2", samples)
+    result = model.evaluate(eval_datasource, value_column="text", record_predictions=True, tags={"test"})
+    assert result is not None
+    predictions = model.predictions(tag="test")
+    assert len(predictions) == 2
+    assert all(p.tags == {"test"} for p in predictions)
+    assert all(p.expected_label == s["label"] for p, s in zip(predictions, samples))
+
+
+def test_predict(model: ClassificationModel, label_names: list[str]):
+    predictions = model.predict(["Do you love soup?", "Are cats cute?"])
+    assert len(predictions) == 2
+    assert predictions[0].label == 0
+    assert predictions[0].label_name == label_names[0]
+    assert 0 <= predictions[0].confidence <= 1
+    assert predictions[1].label == 1
+    assert predictions[1].label_name == label_names[1]
+    assert 0 <= predictions[1].confidence <= 1
+
+
+def test_predict_unauthenticated(unauthenticated, model: ClassificationModel):
+    with pytest.raises(ValueError, match="Invalid API key"):
+        model.predict(["Do you love soup?", "Are cats cute?"])
+
+
+def test_predict_unauthorized(unauthorized, model: ClassificationModel):
+    with pytest.raises(LookupError):
+        model.predict(["Do you love soup?", "Are cats cute?"])
+
+
+def test_predict_constraint_violation(memoryset: LabeledMemoryset):
+    model = ClassificationModel.create(
+        "test_model_lookup_count_too_high", memoryset, num_classes=2, memory_lookup_count=memoryset.length + 2
+    )
+    with pytest.raises(RuntimeError):
+        model.predict("test")
+
+
+def test_record_prediction_feedback(model: ClassificationModel):
+    predictions = model.predict(["Do you love soup?", "Are cats cute?"])
+    expected_labels = [0, 1]
+    model.record_feedback(
+        {
+            "prediction_id": p.prediction_id,
+            "category": "correct",
+            "value": p.label == expected_label,
+        }
+        for expected_label, p in zip(expected_labels, predictions)
+    )
+
+
+def test_record_prediction_feedback_missing_category(model: ClassificationModel):
+    prediction = model.predict("Do you love soup?")
+    with pytest.raises(ValueError):
+        model.record_feedback({"prediction_id": prediction.prediction_id, "value": True})
+
+
+def test_record_prediction_feedback_invalid_value(model: ClassificationModel):
+    prediction = model.predict("Do you love soup?")
+    with pytest.raises(ValueError, match=r"Invalid input.*"):
+        model.record_feedback({"prediction_id": prediction.prediction_id, "category": "correct", "value": "invalid"})
+
+
+def test_record_prediction_feedback_invalid_prediction_id(model: ClassificationModel):
+    with pytest.raises(ValueError, match=r"Invalid input.*"):
+        model.record_feedback({"prediction_id": "invalid", "category": "correct", "value": True})
+
+
+def test_predict_with_memoryset_override(model: ClassificationModel, hf_dataset: Dataset):
+    inverted_labeled_memoryset = LabeledMemoryset.from_hf_dataset(
+        "test_memoryset_inverted_labels",
+        hf_dataset.map(lambda x: {"label": 1 if x["label"] == 0 else 0}),
+        value_column="text",
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+    )
+    with model.use_memoryset(inverted_labeled_memoryset):
+        predictions = model.predict(["Do you love soup?", "Are cats cute?"])
+        assert predictions[0].label == 1
+        assert predictions[1].label == 0
+
+    predictions = model.predict(["Do you love soup?", "Are cats cute?"])
+    assert predictions[0].label == 0
+    assert predictions[1].label == 1
+
+
+def test_predict_with_expected_labels(model: ClassificationModel):
+    prediction = model.predict("Do you love soup?", expected_labels=1)
+    assert prediction.expected_label == 1
+
+
+def test_predict_with_expected_labels_invalid_input(model: ClassificationModel):
+    # invalid number of expected labels for batch prediction
+    with pytest.raises(ValueError, match=r"Invalid input.*"):
+        model.predict(["Do you love soup?", "Are cats cute?"], expected_labels=[0])
+    # invalid label value
+    with pytest.raises(ValueError):
+        model.predict("Do you love soup?", expected_labels=5)
+
+
+def test_last_prediction_with_batch(model: ClassificationModel):
+    predictions = model.predict(["Do you love soup?", "Are cats cute?"])
+    assert model.last_prediction is not None
+    assert model.last_prediction.prediction_id == predictions[-1].prediction_id
+    assert model.last_prediction.input_value == "Are cats cute?"
+    assert model._last_prediction_was_batch is True
+
+
+def test_last_prediction_with_single(model: ClassificationModel):
+    # Test that last_prediction is updated correctly with single prediction
+    prediction = model.predict("Do you love soup?")
+    assert model.last_prediction is not None
+    assert model.last_prediction.prediction_id == prediction.prediction_id
+    assert model.last_prediction.input_value == "Do you love soup?"
+    assert model._last_prediction_was_batch is False

orca_sdk/conftest.py

@@ -0,0 +1,117 @@
+import logging
+import os
+from typing import Generator
+from uuid import uuid4
+
+import pytest
+from datasets import ClassLabel, Dataset, Features, Value
+
+from ._generated_api_client.client import set_headers
+from ._utils.auth import _create_api_key, _delete_org
+from .classification_model import ClassificationModel
+from .datasource import Datasource
+from .embedding_model import PretrainedEmbeddingModel
+from .labeled_memoryset import LabeledMemoryset
+from .orca_credentials import OrcaCredentials
+
+logging.basicConfig(level=logging.INFO)
+
+os.environ["ORCA_API_URL"] = os.environ.get("ORCA_API_URL", "http://localhost:1584/")
+
+
+def _create_org_id():
+    # UUID start to identify test data (0xtest...)
+    return "10e50000-0000-4000-a000-" + str(uuid4())[24:]
+
+
+@pytest.fixture(scope="session")
+def org_id():
+    return _create_org_id()
+
+
+@pytest.fixture(autouse=True, scope="session")
+def api_key(org_id) -> Generator[str, None, None]:
+    api_key = _create_api_key(org_id=org_id, name="orca_sdk_test")
+    OrcaCredentials.set_api_key(api_key, check_validity=True)
+    yield api_key
+    _delete_org(org_id)
+
+
+@pytest.fixture(autouse=True)
+def authenticated(api_key):
+    OrcaCredentials.set_api_key(api_key, check_validity=False)
+
+
+@pytest.fixture()
+def unauthenticated(api_key):
+    OrcaCredentials.set_api_key(str(uuid4()), check_validity=False)
+    yield
+    # Need to reset the api key to the original api key so following tests don't fail
+    OrcaCredentials.set_api_key(api_key, check_validity=False)
+
+
+@pytest.fixture()
+def other_org_id():
+    return _create_org_id()
+
+
+@pytest.fixture()
+def unauthorized(api_key, other_org_id):
+    different_api_key = _create_api_key(org_id=other_org_id, name="orca_sdk_test_other_org")
+    OrcaCredentials.set_api_key(different_api_key, check_validity=False)
+    yield
+    OrcaCredentials.set_api_key(api_key, check_validity=False)
+    _delete_org(other_org_id)
+
+
+@pytest.fixture(scope="session")
+def label_names():
+    return ["soup", "cats"]
+
+
+SAMPLE_DATA = [
+    {"text": "i love soup", "label": 0, "key": "val1", "score": 0.1, "source_id": "s1"},
+    {"text": "cats are cute", "label": 1, "key": "val2", "score": 0.2, "source_id": "s2"},
+    {"text": "soup is good", "label": 0, "key": "val3", "score": 0.3, "source_id": "s3"},
+    {"text": "i love cats", "label": 1, "key": "val4", "score": 0.4, "source_id": "s4"},
+    {"text": "everyone loves cats", "label": 1, "key": "val5", "score": 0.5, "source_id": "s5"},
+    {"text": "soup is great for the winter", "label": 0, "key": "val6", "score": 0.6, "source_id": "s6"},
+]
+
+
+@pytest.fixture(scope="session")
+def hf_dataset(label_names):
+    return Dataset.from_list(
+        SAMPLE_DATA,
+        features=Features(
+            {
+                "text": Value("string"),
+                "label": ClassLabel(names=label_names),
+                "key": Value("string"),
+                "score": Value("float"),
+                "source_id": Value("string"),
+            }
+        ),
+    )
+
+
+@pytest.fixture(scope="session")
+def datasource(hf_dataset) -> Datasource:
+    return Datasource.from_hf_dataset("test_datasource", hf_dataset)
+
+
+@pytest.fixture(scope="session")
+def memoryset(datasource) -> LabeledMemoryset:
+    return LabeledMemoryset.create(
+        "test_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        value_column="text",
+        source_id_column="source_id",
+        max_seq_length_override=32,
+    )
+
+
+@pytest.fixture(scope="session")
+def model(memoryset) -> ClassificationModel:
+    return ClassificationModel.create("test_model", memoryset, num_classes=2, memory_lookup_count=3)