orca-sdk 0.1.11__py3-none-any.whl → 0.1.12__py3-none-any.whl

orca_sdk/memoryset_test.py

@@ -5,7 +5,6 @@ import pytest
 from datasets.arrow_dataset import Dataset
 
 from .classification_model import ClassificationModel
-from .conftest import skip_in_ci, skip_in_prod
 from .datasource import Datasource
 from .embedding_model import PretrainedEmbeddingModel
 from .memoryset import (
@@ -126,6 +125,33 @@ def test_create_empty_scored_memoryset():
         ScoredMemoryset.drop(name, if_not_exists="ignore")
 
 
+def test_create_empty_partitioned_labeled_memoryset():
+    name = f"test_empty_partitioned_labeled_{uuid4()}"
+    label_names = ["negative", "positive"]
+    try:
+        memoryset = LabeledMemoryset.create(
+            name, label_names=label_names, partitioned=True, description="empty partitioned labeled test"
+        )
+        assert memoryset is not None
+        assert memoryset.name == name
+        assert memoryset.length == 0
+        assert memoryset.partitioned is True
+
+        # inserting with partition_id should work
+        memoryset.insert(dict(value="i love soup", label=1, partition_id="p1"))
+        memoryset.insert(dict(value="cats are cute", label=0, partition_id="p2"))
+        assert memoryset.length == 2
+    finally:
+        LabeledMemoryset.drop(name, if_not_exists="ignore")
+
+
+def test_non_partitioned_memoryset_has_partitioned_false(
+    readonly_partitioned_memoryset: LabeledMemoryset, readonly_memoryset: LabeledMemoryset
+):
+    assert readonly_partitioned_memoryset.partitioned is True
+    assert readonly_memoryset.partitioned is False
+
+
 def test_create_memoryset_unauthenticated(unauthenticated_client, datasource):
     with unauthenticated_client.use():
         with pytest.raises(ValueError, match="Invalid API key"):
@@ -718,7 +744,6 @@ def test_insert_memories(writable_memoryset: LabeledMemoryset):
     assert last_memory.source_id == "test"
 
 
-@skip_in_prod("Production memorysets do not have session consistency guarantees")
 def test_update_memories(writable_memoryset: LabeledMemoryset, hf_dataset: Dataset):
     # We've combined the update tests into one to avoid multiple expensive requests for a writable_memoryset
 
@@ -726,10 +751,9 @@ def test_update_memories(writable_memoryset: LabeledMemoryset, hf_dataset: Datas
     memory_id = writable_memoryset[0].memory_id
     updated_count = writable_memoryset.update(dict(memory_id=memory_id, value="i love soup so much"))
     assert updated_count == 1
-    updated_memory = writable_memoryset.get(memory_id)
+    updated_memory = writable_memoryset.get(memory_id, consistency_level="Strong")
     assert updated_memory.value == "i love soup so much"
     assert updated_memory.label == hf_dataset[0]["label"]
-    assert writable_memoryset.get(memory_id).value == "i love soup so much"
 
     # test updating a memory instance
     memory = writable_memoryset[0]
@@ -748,11 +772,10 @@ def test_update_memories(writable_memoryset: LabeledMemoryset, hf_dataset: Datas
         batch_size=1,
     )
     assert updated_count == 2
-    assert writable_memoryset.get(memory_ids[0]).value == "i love soup so much"
-    assert writable_memoryset.get(memory_ids[1]).value == "cats are so cute"
+    assert writable_memoryset.get(memory_ids[0], consistency_level="Strong").value == "i love soup so much"
+    assert writable_memoryset.get(memory_ids[1], consistency_level="Strong").value == "cats are so cute"
 
 
-@skip_in_prod("Production memorysets do not have session consistency guarantees")
 def test_update_memory_metadata(writable_memoryset: LabeledMemoryset):
     memory = writable_memoryset[0]
     assert memory.metadata["key"] == "g1"
@@ -769,7 +792,7 @@ def test_update_memory_metadata(writable_memoryset: LabeledMemoryset):
 
     # Can explicitly clear metadata by passing metadata={}
     writable_memoryset.update(dict(memory_id=memory.memory_id, metadata={}))
-    updated = writable_memoryset.get(memory.memory_id)
+    updated = writable_memoryset.get(memory.memory_id, consistency_level="Strong")
     assert updated.metadata == {}, "Metadata should be cleared when explicitly set to {}"
 
 
@@ -782,7 +805,9 @@ def test_update_memories_by_filter(writable_memoryset: LabeledMemoryset):
     )
     assert updated_count == 2
     assert len(writable_memoryset) == initial_length
-    updated_memories = writable_memoryset.query(filters=[("source_id", "in", source_ids_to_update)])
+    updated_memories = writable_memoryset.query(
+        filters=[("source_id", "in", source_ids_to_update)], consistency_level="Strong"
+    )
     assert len(updated_memories) == 2
     assert all(memory.label == 1 for memory in updated_memories)
 
@@ -980,15 +1005,16 @@ def test_scored_memoryset(scored_memoryset: ScoredMemoryset):
     assert lookup[0].score < 0.11
 
 
-@skip_in_prod("Production memorysets do not have session consistency guarantees")
 def test_update_scored_memory(scored_memoryset: ScoredMemoryset):
     # we are only updating an inconsequential metadata field so that we don't affect other tests
     memory = scored_memoryset[0]
     assert memory.label == 0
     scored_memoryset.update(dict(memory_id=memory.memory_id, label=3))
-    assert scored_memoryset[0].label == 3
-    memory.update(label=4)
-    assert scored_memoryset[0].label == 4
+    memory = scored_memoryset.get(memory.memory_id, consistency_level="Strong")
+    assert memory.label == 3
+    memory = memory.update(label=4)
+    memory = scored_memoryset.get(memory.memory_id, consistency_level="Strong")
+    assert memory.label == 4
 
 
 @pytest.mark.asyncio

orca_sdk/regression_model.py

@@ -1,25 +1,33 @@
 from __future__ import annotations
 
-import logging
 from contextlib import contextmanager
 from datetime import datetime
-from typing import Any, Generator, Iterable, Literal, cast, overload
-
-from datasets import Dataset
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generator,
+    Iterable,
+    Literal,
+    Sequence,
+    cast,
+    overload,
+)
 
-from ._shared.metrics import RegressionMetrics, calculate_regression_metrics
-from ._utils.common import UNSET, CreateMode, DropMode
+from ._utils.common import UNSET, CreateMode, DropMode, logger
 from .client import (
     ListPredictionsRequest,
     OrcaClient,
     PredictiveModelUpdate,
     RARHeadType,
+)
+from .client import RegressionMetrics as RegressionMetricsResponse
+from .client import (
     RegressionModelMetadata,
     RegressionPredictionRequest,
 )
 from .datasource import Datasource
 from .job import Job
-from .memoryset import ScoredMemoryset
+from .memoryset import ConsistencyLevel, ScoredMemoryset
 from .telemetry import (
     RegressionPrediction,
     TelemetryMode,
@@ -27,7 +35,107 @@ from .telemetry import (
     _parse_feedback,
 )
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    # Peer dependency - user has datasets if they have a Dataset object
+    from datasets import Dataset as HFDataset  # type: ignore
+    from pandas import DataFrame as PandasDataFrame  # type: ignore
+
+
+class RegressionMetrics:
+    """
+    Metrics for evaluating regression model performance.
+
+    Attributes:
+        coverage: Percentage of predictions that are not none
+        mse: Mean squared error of the predictions
+        rmse: Root mean squared error of the predictions
+        mae: Mean absolute error of the predictions
+        r2: R-squared score (coefficient of determination) of the predictions
+        explained_variance: Explained variance score of the predictions
+        loss: Mean squared error loss of the predictions
+        anomaly_score_mean: Mean of anomaly scores across the dataset
+        anomaly_score_median: Median of anomaly scores across the dataset
+        anomaly_score_variance: Variance of anomaly scores across the dataset
+    """
+
+    coverage: float
+    mse: float
+    rmse: float
+    mae: float
+    r2: float
+    explained_variance: float
+    loss: float
+    anomaly_score_mean: float | None
+    anomaly_score_median: float | None
+    anomaly_score_variance: float | None
+
+    def __init__(self, response: RegressionMetricsResponse):
+        self.coverage = response["coverage"]
+        self.mse = response["mse"]
+        self.rmse = response["rmse"]
+        self.mae = response["mae"]
+        self.r2 = response["r2"]
+        self.explained_variance = response["explained_variance"]
+        self.loss = response["loss"]
+        self.anomaly_score_mean = response.get("anomaly_score_mean")
+        self.anomaly_score_median = response.get("anomaly_score_median")
+        self.anomaly_score_variance = response.get("anomaly_score_variance")
+        for warning in response.get("warnings", []):
+            logger.warning(warning)
+
+    def __repr__(self) -> str:
+        return (
+            "RegressionMetrics({\n"
+            + f"    mae: {self.mae:.4f},\n"
+            + f"    rmse: {self.rmse:.4f},\n"
+            + f"    r2: {self.r2:.4f},\n"
+            + (
+                f"    anomaly_score: {self.anomaly_score_mean:.4f} ± {self.anomaly_score_variance:.4f},\n"
+                if self.anomaly_score_mean
+                else ""
+            )
+            + "})"
+        )
+
+    @classmethod
+    def compute(
+        cls,
+        predictions: Sequence[RegressionPrediction],
+    ) -> RegressionMetrics:
+        """
+        Compute regression metrics from a list of predictions.
+
+        Params:
+            predictions: List of RegressionPrediction objects with expected_score set
+
+        Returns:
+            RegressionMetrics with computed metrics
+
+        Raises:
+            ValueError: If any prediction is missing expected_score
+        """
+        if len(predictions) > 100_000:
+            raise ValueError("Too many predictions, maximum is 100,000")
+        if any(p.expected_score is None for p in predictions):
+            raise ValueError("All predictions must have expected_score set")
+        expected_scores = [cast(float, p.expected_score) for p in predictions]
+        predicted_scores = [p.score for p in predictions]
+        anomaly_scores = (
+            None
+            if any(p.anomaly_score is None for p in predictions)
+            else [cast(float, p.anomaly_score) for p in predictions]
+        )
+
+        client = OrcaClient._resolve_client()
+        response = client.POST(
+            "/regression_model/metrics",
+            json={
+                "expected_scores": expected_scores,
+                "predicted_scores": predicted_scores,
+                "anomaly_scores": anomaly_scores,
+            },
+        )
+        return cls(response)
 
 
 class RegressionModel:
@@ -105,7 +213,7 @@ class RegressionModel:
             is raised.
         """
         if self._last_prediction_was_batch:
-            logging.warning(
+            logger.warning(
                 "Last prediction was part of a batch prediction, returning the last prediction from the batch"
             )
         if self._last_prediction is None:
@@ -233,7 +341,7 @@ class RegressionModel:
         try:
             client = OrcaClient._resolve_client()
             client.DELETE("/regression_model/{name_or_id}", params={"name_or_id": name_or_id})
-            logging.info(f"Deleted model {name_or_id}")
+            logger.info(f"Deleted model {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
                 raise
@@ -297,6 +405,7 @@ class RegressionModel:
         ] = "include_global",
         use_gpu: bool = True,
         batch_size: int = 100,
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> RegressionPrediction: ...
 
     @overload
@@ -316,6 +425,7 @@ class RegressionModel:
         ] = "include_global",
         use_gpu: bool = True,
         batch_size: int = 100,
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> list[RegressionPrediction]: ...
 
     # TODO: add filter support
@@ -335,6 +445,7 @@ class RegressionModel:
         ] = "include_global",
         use_gpu: bool = True,
         batch_size: int = 100,
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> RegressionPrediction | list[RegressionPrediction]:
         """
         Make predictions using the regression model.
@@ -383,15 +494,13 @@ class RegressionModel:
         client = OrcaClient._resolve_client()
 
         # Convert to list for batching
-        values = value if isinstance(value, list) else [value]
+        values = [value] if isinstance(value, str) else list(value)
         if isinstance(expected_scores, list) and len(expected_scores) != len(values):
             raise ValueError("Invalid input: \n\texpected_scores 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_scores, list):
-            expected_scores = expected_scores
-        elif expected_scores is not None:
+        if expected_scores is not None and isinstance(expected_scores, (float, int)):
             expected_scores = [float(expected_scores)] * len(values)
 
         predictions: list[RegressionPrediction] = []
@@ -410,6 +519,7 @@ class RegressionModel:
                 "use_lookup_cache": use_lookup_cache,
                 "ignore_unlabeled": ignore_unlabeled,
                 "partition_filter_mode": partition_filter_mode,
+                "consistency_level": consistency_level,
             }
             if partition_filter_mode != "ignore_partitions":
                 request_json["partition_ids"] = (
@@ -426,6 +536,7 @@ class RegressionModel:
             if telemetry_on and any(p["prediction_id"] is None for p in response):
                 raise RuntimeError("Failed to save prediction to database.")
 
+            batch_expected = batch_expected_scores or [None] * len(batch_values)
             predictions.extend(
                 RegressionPrediction(
                     prediction_id=prediction["prediction_id"],
@@ -438,8 +549,9 @@ class RegressionModel:
                     model=self,
                     logits=None,
                     input_value=input_value,
+                    expected_score=exp_score,
                 )
-                for prediction, input_value in zip(response, batch_values)
+                for prediction, input_value, exp_score in zip(response, batch_values, batch_expected)
             )
 
         self._last_prediction_was_batch = isinstance(value, list)
@@ -581,25 +693,14 @@ class RegressionModel:
                 params={"model_name_or_id": self.id, "job_id": response["job_id"]},
             )
             assert res["result"] is not None
-            return RegressionMetrics(
-                coverage=res["result"].get("coverage"),
-                mse=res["result"].get("mse"),
-                rmse=res["result"].get("rmse"),
-                mae=res["result"].get("mae"),
-                r2=res["result"].get("r2"),
-                explained_variance=res["result"].get("explained_variance"),
-                loss=res["result"].get("loss"),
-                anomaly_score_mean=res["result"].get("anomaly_score_mean"),
-                anomaly_score_median=res["result"].get("anomaly_score_median"),
-                anomaly_score_variance=res["result"].get("anomaly_score_variance"),
-            )
+            return RegressionMetrics(res["result"])
 
         job = Job(response["job_id"], get_value)
         return job if background else job.result()
 
-    def _evaluate_dataset(
+    def _evaluate_local(
         self,
-        dataset: Dataset,
+        data: Iterable[dict[str, Any]],
         value_column: str,
         score_column: str,
         record_predictions: bool,
@@ -612,37 +713,42 @@ class RegressionModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
     ) -> RegressionMetrics:
-        if len(dataset) == 0:
-            raise ValueError("Evaluation dataset cannot be empty")
-
-        if any(x is None for x in dataset[score_column]):
-            raise ValueError("Evaluation dataset cannot contain None values in the score column")
-
-        predictions = [
-            prediction
-            for i in range(0, len(dataset), batch_size)
-            for prediction in self.predict(
-                dataset[i : i + batch_size][value_column],
-                expected_scores=dataset[i : i + batch_size][score_column],
-                tags=tags,
-                save_telemetry="sync" if record_predictions else "off",
-                prompt=prompt,
-                ignore_unlabeled=ignore_unlabeled,
-                partition_id=dataset[i : i + batch_size][partition_column] if partition_column else None,
-                partition_filter_mode=partition_filter_mode,
-            )
-        ]
-
-        return calculate_regression_metrics(
-            expected_scores=dataset[score_column],
-            predicted_scores=[p.score for p in predictions],
-            anomaly_scores=[p.anomaly_score for p in predictions],
+        values: list[str] = []
+        expected_scores: list[float] = []
+        partition_ids: list[str | None] | None = [] if partition_column else None
+
+        for sample in data:
+            if len(values) >= 100_000:
+                raise ValueError("Upload a Datasource to evaluate against more than 100,000 samples.")
+            values.append(sample[value_column])
+            expected_score = sample[score_column]
+            if expected_score is None:
+                raise ValueError("Expected score is required for all samples")
+            expected_scores.append(expected_score)
+            if partition_ids is not None and partition_column:
+                partition_ids.append(sample[partition_column])
+
+        if not values:
+            raise ValueError("Evaluation data cannot be empty")
+
+        predictions = self.predict(
+            values,
+            expected_scores=expected_scores,
+            tags=tags,
+            save_telemetry="sync" if record_predictions else "off",
+            prompt=prompt,
+            ignore_unlabeled=ignore_unlabeled,
+            partition_id=partition_ids,
+            partition_filter_mode=partition_filter_mode,
+            batch_size=batch_size,
         )
 
+        return RegressionMetrics.compute(predictions)
+
     @overload
     def evaluate(
         self,
-        data: Datasource | Dataset,
+        data: Datasource,
         *,
         value_column: str = "value",
         score_column: str = "score",
@@ -663,7 +769,7 @@ class RegressionModel:
     @overload
     def evaluate(
         self,
-        data: Datasource | Dataset,
+        data: Datasource | HFDataset | PandasDataFrame | Iterable[dict[str, Any]],
         *,
         value_column: str = "value",
         score_column: str = "score",
@@ -683,7 +789,7 @@ class RegressionModel:
 
     def evaluate(
         self,
-        data: Datasource | Dataset,
+        data: Datasource | HFDataset | PandasDataFrame | Iterable[dict[str, Any]],
         *,
         value_column: str = "value",
         score_column: str = "score",
@@ -703,12 +809,13 @@ class RegressionModel:
         Evaluate the regression model on a given dataset or datasource
 
         Params:
-            data: Dataset or Datasource to evaluate the model on
+            data: the data to evaluate the model on. This can be an Orca [`Datasource`][orca_sdk.datasource.Datasource],
+                a Hugging Face [`Dataset`][datasets.Dataset], a pandas [`DataFrame`][pandas.DataFrame], or an iterable of dictionaries.
             value_column: Name of the column that contains the input values to the model
             score_column: Name of the column containing the expected scores
             record_predictions: Whether to record [`RegressionPrediction`][orca_sdk.telemetry.RegressionPrediction]s for analysis
             tags: Optional tags to add to the recorded [`RegressionPrediction`][orca_sdk.telemetry.RegressionPrediction]s
-            batch_size: Batch size for processing Dataset inputs (only used when input is a Dataset)
+            batch_size: Batch size for processing the data inputs (not used for Datasource inputs)
             prompt: Optional prompt for instruction-tuned embedding models
             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
@@ -752,9 +859,25 @@ class RegressionModel:
                 partition_column=partition_column,
                 partition_filter_mode=partition_filter_mode,
             )
-        elif isinstance(data, Dataset):
-            return self._evaluate_dataset(
-                dataset=data,
+        else:
+            if background:
+                raise ValueError("Background evaluation is only supported for Datasource inputs")
+            try:
+                import pandas as pd  # type: ignore
+
+                if isinstance(data, pd.DataFrame):
+                    data = data.to_dict(orient="records")  # type: ignore
+            except ImportError:
+                pass
+
+            if not hasattr(data, "__iter__"):
+                raise ValueError(
+                    f"Invalid data type: {type(data).__name__}. "
+                    "Expected Iterable[dict], HuggingFace Dataset, or pandas DataFrame."
+                )
+
+            return self._evaluate_local(
+                data=cast(Iterable[dict[str, Any]], data),
                 value_column=value_column,
                 score_column=score_column,
                 record_predictions=record_predictions,
@@ -765,8 +888,6 @@ class RegressionModel:
                 partition_column=partition_column,
                 partition_filter_mode=partition_filter_mode,
             )
-        else:
-            raise ValueError(f"Invalid data type: {type(data)}")
 
     @contextmanager
     def use_memoryset(self, memoryset_override: ScoredMemoryset) -> Generator[None, None, None]:

orca_sdk/regression_model_test.py

@@ -173,10 +173,11 @@ def test_delete_memoryset_with_model_cascade(hf_dataset):
     assert not ScoredMemoryset.exists(memoryset.name)
 
 
-@pytest.mark.parametrize("data_type", ["dataset", "datasource"])
+@pytest.mark.parametrize("data_type", ["dataset", "datasource", "list"])
 def test_evaluate(
     regression_model: RegressionModel,
     eval_datasource: Datasource,
+    eval_data: list[dict],
     eval_dataset: Dataset,
     data_type,
 ):
@@ -184,7 +185,11 @@ def test_evaluate(
     result = (
         regression_model.evaluate(eval_dataset)
         if data_type == "dataset"
-        else regression_model.evaluate(eval_datasource)
+        else (
+            regression_model.evaluate(eval_datasource)
+            if data_type == "datasource"
+            else regression_model.evaluate(eval_data)
+        )
     )
 
     assert isinstance(result, RegressionMetrics)
@@ -365,6 +370,8 @@ def test_predict(regression_model: RegressionModel):
     assert len(predictions) == 2
     assert predictions[0].prediction_id is not None
     assert predictions[1].prediction_id is not None
+    assert predictions[0].score is not None
+    assert predictions[1].score is not None
     assert np.allclose(predictions[0].score, 0.1)
     assert np.allclose(predictions[1].score, 0.9)
     assert 0 <= predictions[0].confidence <= 1
@@ -501,7 +508,7 @@ def test_record_prediction_feedback(regression_model: RegressionModel):
         {
             "prediction_id": p.prediction_id,
             "category": "accurate",
-            "value": abs(p.score - expected_score) < 0.2,
+            "value": abs(p.score - expected_score) < 0.2 if p.score is not None else False,
         }
         for expected_score, p in zip(expected_scores, predictions)
     )
@@ -538,11 +545,19 @@ def test_predict_with_memoryset_override(regression_model: RegressionModel, hf_d
     with regression_model.use_memoryset(inverted_scored_memoryset):
         override_predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
         # With inverted scores, the predictions should be different
+        assert original_predictions[0].score is not None
+        assert original_predictions[1].score is not None
+        assert override_predictions[0].score is not None
+        assert override_predictions[1].score is not None
         assert abs(override_predictions[0].score - original_predictions[0].score) > 0.1
         assert abs(override_predictions[1].score - original_predictions[1].score) > 0.1
 
     # After exiting context, predictions should be back to normal
     new_predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
+    assert new_predictions[0].score is not None
+    assert new_predictions[1].score is not None
+    assert original_predictions[0].score is not None
+    assert original_predictions[1].score is not None
     assert abs(new_predictions[0].score - original_predictions[0].score) < 0.1
     assert abs(new_predictions[1].score - original_predictions[1].score) < 0.1
 

orca_sdk/telemetry.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import logging
 import os
 from abc import ABC
 from datetime import datetime
@@ -8,7 +7,7 @@ from typing import TYPE_CHECKING, Any, Iterable, Literal, Self, overload
 
 from httpx import Timeout
 
-from ._utils.common import UNSET
+from ._utils.common import UNSET, logger
 from .client import (
     LabelPredictionWithMemoriesAndFeedback,
     OrcaClient,
@@ -118,7 +117,7 @@ class FeedbackCategory:
         """
         client = OrcaClient._resolve_client()
         client.DELETE("/telemetry/feedback_category/{name_or_id}", params={"name_or_id": name})
-        logging.info(f"Deleted feedback category {name} with all associated feedback")
+        logger.info(f"Deleted feedback category {name} with all associated feedback")
 
     def __repr__(self):
         return "FeedbackCategory({" + f"name: {self.name}, " + f"value_type: {self.value_type}" + "})"
@@ -175,6 +174,8 @@ class PredictionBase(ABC):
         telemetry: LabelPredictionWithMemoriesAndFeedback | ScorePredictionWithMemoriesAndFeedback | None = None,
         logits: list[float] | None = None,
         input_value: str | None = None,
+        expected_label: int | None = None,
+        expected_score: float | None = None,
     ):
         self.prediction_id = prediction_id
         self.label = label
@@ -187,6 +188,8 @@ class PredictionBase(ABC):
         self.__telemetry = telemetry if telemetry else None
         self.logits = logits
         self._input_value = input_value
+        self._expected_label = expected_label
+        self._expected_score = expected_score
 
     @property
     def _telemetry(self) -> LabelPredictionWithMemoriesAndFeedback | ScorePredictionWithMemoriesAndFeedback:
@@ -537,7 +540,7 @@ class ClassificationPrediction(PredictionBase):
         memoryset: Memoryset that was used to lookup memories to ground the prediction
     """
 
-    label: int
+    label: int | None
     label_name: str
     logits: list[float] | None
     model: ClassificationModel
@@ -562,11 +565,15 @@ class ClassificationPrediction(PredictionBase):
 
     @property
     def expected_label(self) -> int | None:
+        if self._expected_label is not None:
+            return self._expected_label
         assert "label" in self._telemetry
         return self._telemetry["expected_label"]
 
     @property
     def expected_label_name(self) -> str | None:
+        if self._expected_label is not None:
+            return self.memoryset.label_names[self._expected_label]
         assert "label" in self._telemetry
         return self._telemetry["expected_label_name"]
 
@@ -692,14 +699,14 @@ class RegressionPrediction(PredictionBase):
         memoryset: Memoryset that was used to lookup memories to ground the prediction
     """
 
-    score: float
+    score: float | None
     model: RegressionModel
     memoryset: ScoredMemoryset
 
     def __repr__(self):
         return (
             "RegressionPrediction({"
-            + f"score: {self.score:.2f}, "
+            + (f"score: {self.score:.2f}, " if self.score is not None else "score: None, ")
             + f"confidence: {self.confidence:.2f}, "
             + (f"anomaly_score: {self.anomaly_score:.2f}, " if self.anomaly_score is not None else "")
             + f"input_value: '{str(self.input_value)[:100] + '...' if len(str(self.input_value)) > 100 else self.input_value}'"
@@ -720,6 +727,8 @@ class RegressionPrediction(PredictionBase):
 
     @property
     def expected_score(self) -> float | None:
+        if self._expected_score is not None:
+            return self._expected_score
         assert "score" in self._telemetry
         return self._telemetry["expected_score"]