orca-sdk 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

orca_sdk/_shared/metrics_test.py

@@ -0,0 +1,273 @@
+"""
+IMPORTANT:
+- This is a shared file between OrcaLib and the OrcaSDK.
+- Please ensure that it does not have any dependencies on the OrcaLib code.
+- Make sure to edit this file in orcalib/shared and NOT in orca_sdk, since it will be overwritten there.
+"""
+
+from typing import Literal
+
+import numpy as np
+import pytest
+import sklearn.metrics
+
+from .metrics import (
+    calculate_classification_metrics,
+    calculate_pr_curve,
+    calculate_regression_metrics,
+    calculate_roc_curve,
+    softmax,
+)
+
+
+def test_binary_metrics():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.3, 0.2])
+
+    metrics = calculate_classification_metrics(y_true, y_score)
+
+    assert metrics.accuracy == 0.8
+    assert metrics.f1_score == 0.8
+    assert metrics.roc_auc is not None
+    assert metrics.roc_auc > 0.8
+    assert metrics.roc_auc < 1.0
+    assert metrics.pr_auc is not None
+    assert metrics.pr_auc > 0.8
+    assert metrics.pr_auc < 1.0
+    assert metrics.loss is not None
+    assert metrics.loss > 0.0
+
+
+def test_multiclass_metrics_with_2_classes():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([[0.9, 0.1], [0.1, 0.9], [0.2, 0.8], [0.7, 0.3], [0.8, 0.2]])
+
+    metrics = calculate_classification_metrics(y_true, y_score)
+
+    assert metrics.accuracy == 0.8
+    assert metrics.f1_score == 0.8
+    assert metrics.roc_auc is not None
+    assert metrics.roc_auc > 0.8
+    assert metrics.roc_auc < 1.0
+    assert metrics.pr_auc is not None
+    assert metrics.pr_auc > 0.8
+    assert metrics.pr_auc < 1.0
+    assert metrics.loss is not None
+    assert metrics.loss > 0.0
+
+
+@pytest.mark.parametrize(
+    "average, multiclass",
+    [("micro", "ovr"), ("macro", "ovr"), ("weighted", "ovr"), ("micro", "ovo"), ("macro", "ovo"), ("weighted", "ovo")],
+)
+def test_multiclass_metrics_with_3_classes(
+    average: Literal["micro", "macro", "weighted"], multiclass: Literal["ovr", "ovo"]
+):
+    y_true = np.array([0, 1, 1, 0, 2])
+    y_score = np.array([[0.9, 0.1, 0.0], [0.1, 0.9, 0.0], [0.2, 0.8, 0.0], [0.7, 0.3, 0.0], [0.0, 0.0, 1.0]])
+
+    metrics = calculate_classification_metrics(y_true, y_score, average=average, multi_class=multiclass)
+
+    assert metrics.accuracy == 1.0
+    assert metrics.f1_score == 1.0
+    assert metrics.roc_auc is not None
+    assert metrics.roc_auc > 0.8
+    assert metrics.pr_auc is None
+    assert metrics.loss is not None
+    assert metrics.loss > 0.0
+
+
+def test_does_not_modify_logits_unless_necessary():
+    logits = np.array([[0.1, 0.9], [0.2, 0.8], [0.7, 0.3], [0.8, 0.2]])
+    expected_labels = [0, 1, 0, 1]
+    assert calculate_classification_metrics(expected_labels, logits).loss == sklearn.metrics.log_loss(
+        expected_labels, logits
+    )
+
+
+def test_normalizes_logits_if_necessary():
+    logits = np.array([[1.2, 3.9], [1.2, 5.8], [1.2, 2.7], [1.2, 1.3]])
+    expected_labels = [0, 1, 0, 1]
+    assert calculate_classification_metrics(expected_labels, logits).loss == sklearn.metrics.log_loss(
+        expected_labels, logits / logits.sum(axis=1, keepdims=True)
+    )
+
+
+def test_softmaxes_logits_if_necessary():
+    logits = np.array([[-1.2, 3.9], [1.2, -5.8], [1.2, 2.7], [1.2, 1.3]])
+    expected_labels = [0, 1, 0, 1]
+    assert calculate_classification_metrics(expected_labels, logits).loss == sklearn.metrics.log_loss(
+        expected_labels, softmax(logits)
+    )
+
+
+def test_handles_nan_logits():
+    logits = np.array([[np.nan, np.nan], [np.nan, np.nan], [0.1, 0.9], [0.2, 0.8]])
+    expected_labels = [0, 1, 0, 1]
+    metrics = calculate_classification_metrics(expected_labels, logits)
+    assert metrics.loss is None
+    assert metrics.accuracy == 0.25
+    assert metrics.f1_score == 0.25
+    assert metrics.roc_auc is None
+    assert metrics.pr_auc is None
+    assert metrics.pr_curve is None
+    assert metrics.roc_curve is None
+    assert metrics.coverage == 0.5
+
+
+def test_precision_recall_curve():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    pr_curve = calculate_pr_curve(y_true, y_score)
+
+    assert len(pr_curve["precisions"]) == len(pr_curve["recalls"]) == len(pr_curve["thresholds"]) == 6
+    assert np.allclose(pr_curve["precisions"][0], 0.6)
+    assert np.allclose(pr_curve["recalls"][0], 1.0)
+    assert np.allclose(pr_curve["precisions"][-1], 1.0)
+    assert np.allclose(pr_curve["recalls"][-1], 0.0)
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(pr_curve["thresholds"]) >= 0)
+
+
+def test_roc_curve():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    roc_curve = calculate_roc_curve(y_true, y_score)
+
+    assert (
+        len(roc_curve["false_positive_rates"])
+        == len(roc_curve["true_positive_rates"])
+        == len(roc_curve["thresholds"])
+        == 6
+    )
+    assert roc_curve["false_positive_rates"][0] == 1.0
+    assert roc_curve["true_positive_rates"][0] == 1.0
+    assert roc_curve["false_positive_rates"][-1] == 0.0
+    assert roc_curve["true_positive_rates"][-1] == 0.0
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(roc_curve["thresholds"]) >= 0)
+
+
+def test_log_loss_handles_missing_classes_in_y_true():
+    # y_true contains only a subset of classes, but predictions include an extra class column
+    y_true = np.array([0, 1, 0, 1])
+    y_score = np.array(
+        [
+            [0.7, 0.2, 0.1],
+            [0.1, 0.8, 0.1],
+            [0.6, 0.3, 0.1],
+            [0.2, 0.7, 0.1],
+        ]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score)
+    expected_loss = sklearn.metrics.log_loss(y_true, y_score, labels=[0, 1, 2])
+
+    assert metrics.loss is not None
+    assert np.allclose(metrics.loss, expected_loss)
+
+
+def test_precision_recall_curve_max_length():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    pr_curve = calculate_pr_curve(y_true, y_score, max_length=5)
+    assert len(pr_curve["precisions"]) == len(pr_curve["recalls"]) == len(pr_curve["thresholds"]) == 5
+
+    assert np.allclose(pr_curve["precisions"][0], 0.6)
+    assert np.allclose(pr_curve["recalls"][0], 1.0)
+    assert np.allclose(pr_curve["precisions"][-1], 1.0)
+    assert np.allclose(pr_curve["recalls"][-1], 0.0)
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(pr_curve["thresholds"]) >= 0)
+
+
+def test_roc_curve_max_length():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.6, 0.2])
+
+    roc_curve = calculate_roc_curve(y_true, y_score, max_length=5)
+    assert (
+        len(roc_curve["false_positive_rates"])
+        == len(roc_curve["true_positive_rates"])
+        == len(roc_curve["thresholds"])
+        == 5
+    )
+    assert np.allclose(roc_curve["false_positive_rates"][0], 1.0)
+    assert np.allclose(roc_curve["true_positive_rates"][0], 1.0)
+    assert np.allclose(roc_curve["false_positive_rates"][-1], 0.0)
+    assert np.allclose(roc_curve["true_positive_rates"][-1], 0.0)
+
+    # test that thresholds are sorted
+    assert np.all(np.diff(roc_curve["thresholds"]) >= 0)
+
+
+# Regression Metrics Tests
+def test_perfect_regression_predictions():
+    y_true = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
+    y_pred = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
+
+    metrics = calculate_regression_metrics(y_true, y_pred)
+
+    assert metrics.mse == 0.0
+    assert metrics.rmse == 0.0
+    assert metrics.mae == 0.0
+    assert metrics.r2 == 1.0
+    assert metrics.explained_variance == 1.0
+    assert metrics.loss == 0.0
+    assert metrics.anomaly_score_mean is None
+    assert metrics.anomaly_score_median is None
+    assert metrics.anomaly_score_variance is None
+
+
+def test_basic_regression_metrics():
+    y_true = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
+    y_pred = np.array([1.1, 1.9, 3.2, 3.8, 5.1], dtype=np.float32)
+
+    metrics = calculate_regression_metrics(y_true, y_pred)
+
+    # Check that all metrics are reasonable
+    assert metrics.mse > 0.0
+    assert metrics.rmse == pytest.approx(np.sqrt(metrics.mse))
+    assert metrics.mae > 0.0
+    assert 0.0 <= metrics.r2 <= 1.0
+    assert 0.0 <= metrics.explained_variance <= 1.0
+    assert metrics.loss == metrics.mse
+
+    # Check specific values based on the data
+    expected_mse = np.mean((y_true - y_pred) ** 2)
+    assert metrics.mse == pytest.approx(expected_mse)
+
+    expected_mae = np.mean(np.abs(y_true - y_pred))
+    assert metrics.mae == pytest.approx(expected_mae)
+
+
+def test_regression_metrics_with_anomaly_scores():
+    y_true = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
+    y_pred = np.array([1.1, 1.9, 3.2, 3.8, 5.1], dtype=np.float32)
+    anomaly_scores = [0.1, 0.2, 0.15, 0.3, 0.25]
+
+    metrics = calculate_regression_metrics(y_true, y_pred, anomaly_scores)
+
+    assert metrics.anomaly_score_mean == pytest.approx(np.mean(anomaly_scores))
+    assert metrics.anomaly_score_median == pytest.approx(np.median(anomaly_scores))
+    assert metrics.anomaly_score_variance == pytest.approx(np.var(anomaly_scores))
+
+
+def test_regression_metrics_handles_nans():
+    y_true = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    y_pred = np.array([1.1, 1.9, np.nan], dtype=np.float32)
+
+    metrics = calculate_regression_metrics(y_true, y_pred)
+
+    assert np.allclose(metrics.coverage, 0.6666666666666666)
+    assert metrics.mse > 0.0
+    assert metrics.rmse > 0.0
+    assert metrics.mae > 0.0
+    assert 0.0 <= metrics.r2 <= 1.0
+    assert 0.0 <= metrics.explained_variance <= 1.0

orca_sdk/_utils/analysis_ui.py

@@ -28,9 +28,7 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
         filters=[("metrics.neighbor_predicted_label_matches_current_label", "==", False)]
     )
     # Sort memories by confidence score (higher confidence first)
-    suggested_relabels.sort(
-        key=lambda x: (x.metrics and x.metrics.neighbor_predicted_label_confidence) or 0.0, reverse=True
-    )
+    suggested_relabels.sort(key=lambda x: (x.metrics.get("neighbor_predicted_label_confidence", 0.0)), reverse=True)
 
     def update_approved(memory_id: str, selected: bool, current_memory_relabel_map: dict[str, RelabelStatus]):
         current_memory_relabel_map[memory_id]["approved"] = selected
@@ -72,9 +70,9 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
             current_memory_relabel_map[mem_id]["new_label"] = new_label
             confidence = "--"
             current_metrics = current_memory_relabel_map[mem_id]["full_memory"].metrics
-            if current_metrics and new_label == current_metrics.neighbor_predicted_label:
+            if current_metrics and new_label == current_metrics.get("neighbor_predicted_label"):
                 confidence = (
-                    round(current_metrics.neighbor_predicted_label_confidence or 0.0, 2) if current_metrics else 0
+                    round(current_metrics.get("neighbor_predicted_label_confidence", 0.0), 2) if current_metrics else 0
                 )
             return (
                 gr.HTML(
@@ -101,8 +99,8 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
                     memory_id=mem.memory_id,
                     approved=False,
                     new_label=(
-                        mem.metrics.neighbor_predicted_label
-                        if (mem.metrics and isinstance(mem.metrics.neighbor_predicted_label, int))
+                        mem.metrics.get("neighbor_predicted_label")
+                        if (mem.metrics and isinstance(mem.metrics.get("neighbor_predicted_label"), int))
                         else None
                     ),
                     full_memory=mem,
@@ -150,7 +148,11 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
                         )
                 for i, memory_relabel in enumerate(current_memory_relabel_map.values()):
                     mem = memory_relabel["full_memory"]
-                    with gr.Row(equal_height=True, variant="panel", elem_classes="white" if i % 2 == 0 else None):
+                    predicted_label = mem.metrics["neighbor_predicted_label"]
+                    predicted_label_name = label_names[predicted_label]
+                    predicted_label_confidence = mem.metrics.get("neighbor_predicted_label_confidence", 0)
+
+                    with gr.Row(equal_height=True, variant="panel"):
                         with gr.Column(scale=9):
                             assert isinstance(mem.value, str)
                             gr.Markdown(mem.value, label="Value", height=50)
@@ -160,12 +162,12 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
                             dropdown = gr.Dropdown(
                                 choices=[f"{label_name} ({i})" for i, label_name in enumerate(label_names)],
                                 label="SuggestedLabel",
-                                value=f"{label_names[mem.metrics.neighbor_predicted_label]} ({mem.metrics.neighbor_predicted_label})",
+                                value=f"{predicted_label_name} ({predicted_label})",
                                 interactive=True,
                                 container=False,
                             )
                             confidence = gr.HTML(
-                                f"<p style='font-size: 10px; color: #888;'>Confidence: {mem.metrics.neighbor_predicted_label_confidence:.2f}</p>",
+                                f"<p style='font-size: 10px; color: #888;'>Confidence: {predicted_label_confidence:.2f}</p>",
                                 elem_classes="no-padding",
                             )
                             dropdown.change(

orca_sdk/_utils/analysis_ui_style.css

@@ -1,6 +1,3 @@
-.white {
-  background-color: white;
-}
 .centered input {
   margin: auto;
 }

orca_sdk/_utils/auth.py

@@ -2,61 +2,59 @@
 
 import logging
 import os
-from typing import List
+from typing import List, Literal
 
 from dotenv import load_dotenv
 
-from .._generated_api_client.api import (
-    check_authentication,
-    create_api_key,
-    delete_api_key,
-    delete_org,
-    list_api_keys,
-)
-from .._generated_api_client.client import headers_context, set_base_url, set_headers
-from .._generated_api_client.models import CreateApiKeyRequest
-from .._generated_api_client.models.api_key_metadata import ApiKeyMetadata
+from ..client import ApiKeyMetadata, orca_api
+from ..credentials import OrcaCredentials
 from .common import DropMode
 
 load_dotenv()  # this needs to be here to ensure env is populated before accessing it
+
+# the defaults here must match nautilus and lighthouse config defaults
 _ORCA_ROOT_ACCESS_API_KEY = os.environ.get("ORCA_ROOT_ACCESS_API_KEY", "00000000-0000-0000-0000-000000000000")
+_DEFAULT_ORG_ID = os.environ.get("DEFAULT_ORG_ID", "10e50000-0000-4000-a000-a78dca14af3a")
 
 
-def _create_api_key(org_id: str, name: str) -> str:
+def _create_api_key(org_id: str, name: str, scopes: list[Literal["ADMINISTER", "PREDICT"]] = ["ADMINISTER"]) -> str:
     """Creates an API key for the given organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        res = create_api_key(body=CreateApiKeyRequest(name=name))
-        return res.api_key
+    response = orca_api.POST(
+        "/auth/api_key",
+        json={"name": name, "scope": scopes},
+        headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id},
+    )
+    return response["api_key"]
 
 
 def _list_api_keys(org_id: str) -> List[ApiKeyMetadata]:
     """Lists all API keys for the given organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        return list_api_keys()
+    return orca_api.GET("/auth/api_key", headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id})
 
 
 def _delete_api_key(org_id: str, name: str, if_not_exists: DropMode = "error") -> None:
     """Deletes the API key with the given name from the organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        try:
-            delete_api_key(name_or_id=name)
-        except LookupError:
-            if if_not_exists == "error":
-                raise
+    try:
+        orca_api.DELETE(
+            "/auth/api_key/{name_or_id}",
+            params={"name_or_id": name},
+            headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id},
+        )
+    except LookupError:
+        if if_not_exists == "error":
+            raise
 
 
 def _delete_org(org_id: str) -> None:
     """Deletes the organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        delete_org()
+    orca_api.DELETE("/auth/org", headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id})
 
 
-def _authenticate_local_api(org_id: str = "10e50000-0000-4000-a000-a78dca14af3a", api_key_name: str = "local") -> None:
+def _authenticate_local_api(org_id: str = _DEFAULT_ORG_ID, api_key_name: str = "local") -> None:
     """Connect to the local API at http://localhost:1584/ and authenticate with a new API key"""
-    set_base_url("http://localhost:1584/")
     _delete_api_key(org_id, api_key_name, if_not_exists="ignore")
-    set_headers({"Api-Key": _create_api_key(org_id, api_key_name)})
-    check_authentication()
+    OrcaCredentials.set_api_url("http://localhost:1584")
+    OrcaCredentials.set_api_key(_create_api_key(org_id, api_key_name))
     logging.info(f"Authenticated against local API at 'http://localhost:1584' with '{api_key_name}' API key")
 
 

orca_sdk/_utils/data_parsing.py

@@ -1,6 +1,7 @@
 import pickle
 from dataclasses import asdict, is_dataclass
 from os import PathLike
+from tempfile import TemporaryDirectory
 from typing import Any, cast
 
 from datasets import Dataset
@@ -40,7 +41,24 @@ def parse_batch(batch: Any, column_names: list[str] | None = None) -> list[dict]
     return [{key: batch[key][idx] for key in keys} for idx in range(batch_size)]
 
 
-def hf_dataset_from_torch(torch_data: TorchDataLoader | TorchDataset, column_names: list[str] | None = None) -> Dataset:
+def hf_dataset_from_torch(
+    torch_data: TorchDataLoader | TorchDataset, column_names: list[str] | None = None, ignore_cache=False
+) -> Dataset:
+    """
+    Create a HuggingFace Dataset from a PyTorch DataLoader or Dataset.
+
+    NOTE:  It's important to ignore the cached files when testing (i.e., ignore_cache=Ture), because
+    cached results can ignore changes you've made to tests. This can make a test appear to succeed
+    when it's actually broken or vice versa.
+
+    Params:
+        torch_data: A PyTorch DataLoader or Dataset object to create the HuggingFace Dataset from.
+        column_names: Optional list of column names to use for the dataset. If not provided,
+            the column names will be inferred from the data.
+        ignore_cache: If True, the dataset will not be cached on disk.
+    Returns:
+        A HuggingFace Dataset object containing the data from the PyTorch DataLoader or Dataset.
+    """
     if isinstance(torch_data, TorchDataLoader):
         dataloader = torch_data
     else:
@@ -50,7 +68,15 @@ def hf_dataset_from_torch(torch_data: TorchDataLoader | TorchDataset, column_nam
         for batch in dataloader:
             yield from parse_batch(batch, column_names=column_names)
 
-    return cast(Dataset, Dataset.from_generator(generator))
+    if ignore_cache:
+        with TemporaryDirectory() as temp_dir:
+            ds = Dataset.from_generator(generator, cache_dir=temp_dir)
+    else:
+        ds = Dataset.from_generator(generator)
+
+    if not isinstance(ds, Dataset):
+        raise ValueError(f"Failed to create dataset from generator: {type(ds)}")
+    return ds
 
 
 def hf_dataset_from_disk(file_path: str | PathLike) -> Dataset:

orca_sdk/_utils/data_parsing_test.py

@@ -29,11 +29,11 @@ class PytorchDictDataset(TorchDataset):
 def test_hf_dataset_from_torch_dict():
     # Given a Pytorch dataset that returns a dictionary for each item
     dataset = PytorchDictDataset()
-    hf_dataset = hf_dataset_from_torch(dataset)
+    hf_dataset = hf_dataset_from_torch(dataset, ignore_cache=True)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
-    assert set(hf_dataset.column_names) == {"text", "label", "key", "score", "source_id"}
+    assert set(hf_dataset.column_names) == {"value", "label", "key", "score", "source_id"}
 
 
 class PytorchTupleDataset(TorchDataset):
@@ -41,7 +41,7 @@ class PytorchTupleDataset(TorchDataset):
         self.data = SAMPLE_DATA
 
     def __getitem__(self, i):
-        return self.data[i]["text"], self.data[i]["label"]
+        return self.data[i]["value"], self.data[i]["label"]
 
     def __len__(self):
         return len(self.data)
@@ -51,11 +51,11 @@ def test_hf_dataset_from_torch_tuple():
     # Given a Pytorch dataset that returns a tuple for each item
     dataset = PytorchTupleDataset()
     # And the correct number of column names passed in
-    hf_dataset = hf_dataset_from_torch(dataset, column_names=["text", "label"])
+    hf_dataset = hf_dataset_from_torch(dataset, column_names=["value", "label"], ignore_cache=True)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
-    assert hf_dataset.column_names == ["text", "label"]
+    assert hf_dataset.column_names == ["value", "label"]
 
 
 def test_hf_dataset_from_torch_tuple_error():
@@ -63,7 +63,7 @@ def test_hf_dataset_from_torch_tuple_error():
     dataset = PytorchTupleDataset()
     # Then the HF dataset should raise an error if no column names are passed in
     with pytest.raises(DatasetGenerationError):
-        hf_dataset_from_torch(dataset)
+        hf_dataset_from_torch(dataset, ignore_cache=True)
 
 
 def test_hf_dataset_from_torch_tuple_error_not_enough_columns():
@@ -71,7 +71,7 @@ def test_hf_dataset_from_torch_tuple_error_not_enough_columns():
     dataset = PytorchTupleDataset()
     # Then the HF dataset should raise an error if not enough column names are passed in
     with pytest.raises(DatasetGenerationError):
-        hf_dataset_from_torch(dataset, column_names=["value"])
+        hf_dataset_from_torch(dataset, column_names=["value"], ignore_cache=True)
 
 
 DatasetTuple = namedtuple("DatasetTuple", ["value", "label"])
@@ -82,7 +82,7 @@ class PytorchNamedTupleDataset(TorchDataset):
         self.data = SAMPLE_DATA
 
     def __getitem__(self, i):
-        return DatasetTuple(self.data[i]["text"], self.data[i]["label"])
+        return DatasetTuple(self.data[i]["value"], self.data[i]["label"])
 
     def __len__(self):
         return len(self.data)
@@ -92,7 +92,7 @@ def test_hf_dataset_from_torch_named_tuple():
     # Given a Pytorch dataset that returns a namedtuple for each item
     dataset = PytorchNamedTupleDataset()
     # And no column names are passed in
-    hf_dataset = hf_dataset_from_torch(dataset)
+    hf_dataset = hf_dataset_from_torch(dataset, ignore_cache=True)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
@@ -110,7 +110,7 @@ class PytorchDataclassDataset(TorchDataset):
         self.data = SAMPLE_DATA
 
     def __getitem__(self, i):
-        return DatasetItem(text=self.data[i]["text"], label=self.data[i]["label"])
+        return DatasetItem(text=self.data[i]["value"], label=self.data[i]["label"])
 
     def __len__(self):
         return len(self.data)
@@ -119,7 +119,7 @@ class PytorchDataclassDataset(TorchDataset):
 def test_hf_dataset_from_torch_dataclass():
     # Given a Pytorch dataset that returns a dataclass for each item
     dataset = PytorchDataclassDataset()
-    hf_dataset = hf_dataset_from_torch(dataset)
+    hf_dataset = hf_dataset_from_torch(dataset, ignore_cache=True)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
@@ -131,7 +131,7 @@ class PytorchInvalidDataset(TorchDataset):
         self.data = SAMPLE_DATA
 
     def __getitem__(self, i):
-        return [self.data[i]["text"], self.data[i]["label"]]
+        return [self.data[i]["value"], self.data[i]["label"]]
 
     def __len__(self):
         return len(self.data)
@@ -142,7 +142,7 @@ def test_hf_dataset_from_torch_invalid_dataset():
     dataset = PytorchInvalidDataset()
     # Then the HF dataset should raise an error
     with pytest.raises(DatasetGenerationError):
-        hf_dataset_from_torch(dataset)
+        hf_dataset_from_torch(dataset, ignore_cache=True)
 
 
 def test_hf_dataset_from_torchdataloader():
@@ -150,10 +150,10 @@ def test_hf_dataset_from_torchdataloader():
     dataset = PytorchDictDataset()
 
     def collate_fn(x: list[dict]):
-        return {"value": [item["text"] for item in x], "label": [item["label"] for item in x]}
+        return {"value": [item["value"] for item in x], "label": [item["label"] for item in x]}
 
     dataloader = TorchDataLoader(dataset, batch_size=3, collate_fn=collate_fn)
-    hf_dataset = hf_dataset_from_torch(dataloader)
+    hf_dataset = hf_dataset_from_torch(dataloader, ignore_cache=True)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)