hirundo 0.1.16__py3-none-any.whl → 0.1.21__py3-none-any.whl

hirundo/{dataset_optimization.py → dataset_qa.py}

@@ -1,25 +1,25 @@
 import datetime
 import json
 import typing
-from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Generator
 from enum import Enum
 from typing import overload
 
 import httpx
-import requests
 from pydantic import BaseModel, Field, model_validator
 from tqdm import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
 
-from hirundo._constraints import HirundoUrl
+from hirundo._constraints import validate_labeling_info, validate_url
 from hirundo._env import API_HOST
 from hirundo._headers import get_headers
-from hirundo._http import raise_for_status_with_reason
+from hirundo._http import raise_for_status_with_reason, requests
 from hirundo._iter_sse_retrying import aiter_sse_retrying, iter_sse_retrying
 from hirundo._timeouts import MODIFY_TIMEOUT, READ_TIMEOUT
+from hirundo._urls import HirundoUrl
 from hirundo.dataset_enum import DatasetMetadataType, LabelingType
-from hirundo.dataset_optimization_results import DatasetOptimizationResults
+from hirundo.dataset_qa_results import DatasetQAResults
+from hirundo.labeling import YOLO, LabelingInfo
 from hirundo.logger import get_logger
 from hirundo.storage import ResponseStorageConfig, StorageConfig
 from hirundo.unzip import download_and_extract_zip
@@ -29,7 +29,7 @@ logger = get_logger(__name__)
 
 class HirundoError(Exception):
     """
-    Custom exception used to indicate errors in `hirundo` dataset optimization runs
+    Custom exception used to indicate errors in `hirundo` dataset QA runs
     """
 
     pass
@@ -50,14 +50,14 @@ class RunStatus(Enum):
 
 
 STATUS_TO_TEXT_MAP = {
-    RunStatus.STARTED.value: "Optimization run in progress. Downloading dataset",
-    RunStatus.PENDING.value: "Optimization run queued and not yet started",
-    RunStatus.SUCCESS.value: "Optimization run completed successfully",
-    RunStatus.FAILURE.value: "Optimization run failed",
+    RunStatus.STARTED.value: "Dataset QA run in progress. Downloading dataset",
+    RunStatus.PENDING.value: "Dataset QA run queued and not yet started",
+    RunStatus.SUCCESS.value: "Dataset QA run completed successfully",
+    RunStatus.FAILURE.value: "Dataset QA run failed",
     RunStatus.AWAITING_MANUAL_APPROVAL.value: "Awaiting manual approval",
-    RunStatus.RETRY.value: "Optimization run failed. Retrying",
-    RunStatus.REVOKED.value: "Optimization run was cancelled",
-    RunStatus.REJECTED.value: "Optimization run was rejected",
+    RunStatus.RETRY.value: "Dataset QA run failed. Retrying",
+    RunStatus.REVOKED.value: "Dataset QA run was cancelled",
+    RunStatus.REJECTED.value: "Dataset QA run was rejected",
 }
 STATUS_TO_PROGRESS_MAP = {
     RunStatus.STARTED.value: 0.0,
@@ -71,99 +71,51 @@ STATUS_TO_PROGRESS_MAP = {
 }
 
 
-class Metadata(BaseModel, ABC):
-    type: DatasetMetadataType
-
-    @property
-    @abstractmethod
-    def metadata_url(self) -> HirundoUrl:
-        raise NotImplementedError()
-
-
-class HirundoCSV(Metadata):
+class ClassificationRunArgs(BaseModel):
+    image_size: typing.Optional[tuple[int, int]] = (224, 224)
     """
-    A dataset metadata file in the Hirundo CSV format
+    Size (width, height) to which to resize classification images.
+    It is recommended to keep this value at (224, 224) unless your classes are differentiated by very small differences.
     """
-
-    type: DatasetMetadataType = DatasetMetadataType.HIRUNDO_CSV
-    csv_url: HirundoUrl
+    upsample: typing.Optional[bool] = False
     """
-    The URL to access the dataset metadata CSV file.
-    e.g. `s3://my-bucket-name/my-folder/my-metadata.csv`, `gs://my-bucket-name/my-folder/my-metadata.csv`,
-    or `ssh://my-username@my-repo-name/my-folder/my-metadata.csv`
-    (or `file:///datasets/my-folder/my-metadata.csv` if using LOCAL storage type with on-premises installation)
+    Whether to upsample the dataset to attempt to balance the classes.
     """
 
-    @property
-    def metadata_url(self) -> HirundoUrl:
-        return self.csv_url
-
 
-class COCO(Metadata):
+class ObjectDetectionRunArgs(ClassificationRunArgs):
+    min_abs_bbox_size: typing.Optional[int] = None
     """
-    A dataset metadata file in the COCO format
+    Minimum valid size (in pixels) of a bounding box to keep it in the dataset for QA.
     """
-
-    type: DatasetMetadataType = DatasetMetadataType.COCO
-    json_url: HirundoUrl
-    """
-    The URL to access the dataset metadata JSON file.
-    e.g. `s3://my-bucket-name/my-folder/my-metadata.json`, `gs://my-bucket-name/my-folder/my-metadata.json`,
-    or `ssh://my-username@my-repo-name/my-folder/my-metadata.json`
-    (or `file:///datasets/my-folder/my-metadata.json` if using LOCAL storage type with on-premises installation)
-    """
-
-    @property
-    def metadata_url(self) -> HirundoUrl:
-        return self.json_url
-
-
-class YOLO(Metadata):
-    type: DatasetMetadataType = DatasetMetadataType.YOLO
-    data_yaml_url: typing.Optional[HirundoUrl] = None
-    labels_dir_url: HirundoUrl
-
-    @property
-    def metadata_url(self) -> HirundoUrl:
-        return self.labels_dir_url
-
-
-LabelingInfo = typing.Union[HirundoCSV, COCO, YOLO]
-"""
-The dataset labeling info. The dataset labeling info can be one of the following:
-- `DatasetMetadataType.HirundoCSV`: Indicates that the dataset metadata file is a CSV file with the Hirundo format
-
-Currently no other formats are supported. Future versions of `hirundo` may support additional formats.
-"""
-
-
-class VisionRunArgs(BaseModel):
-    upsample: bool = False
+    min_abs_bbox_area: typing.Optional[int] = None
     """
-    Whether to upsample the dataset to attempt to balance the classes.
+    Minimum valid absolute area (in pixels²) of a bounding box to keep it in the dataset for QA.
     """
-    min_abs_bbox_size: int = 0
+    min_rel_bbox_size: typing.Optional[float] = None
     """
-    Minimum valid size (in pixels) of a bounding box to keep it in the dataset for optimization.
+    Minimum valid size (as a fraction of both image height and width) for a bounding box
+    to keep it in the dataset for QA, relative to the corresponding dimension size,
+    i.e. if the bounding box is 10% of the image width and 5% of the image height, it will be kept if this value is 0.05, but not if the
+    value is 0.06 (since both width and height are checked).
     """
-    min_abs_bbox_area: int = 0
+    min_rel_bbox_area: typing.Optional[float] = None
     """
-    Minimum valid absolute area (in pixels²) of a bounding box to keep it in the dataset for optimization.
+    Minimum valid relative area (as a fraction of the image area) of a bounding box to keep it in the dataset for QA.
     """
-    min_rel_bbox_size: float = 0.0
+    crop_ratio: typing.Optional[float] = None
     """
-    Minimum valid size (as a fraction of both image height and width) for a bounding box
-    to keep it in the dataset for optimization, relative to the corresponding dimension size,
-    i.e. if the bounding box is 10% of the image width and 5% of the image height, it will be kept if this value is 0.05, but not if the
-    value is 0.06 (since both width and height are checked).
+    Ratio of the bounding box to crop.
+    Change this value at your own risk. It is recommended to keep it at 1.0 unless you know what you are doing.
     """
-    min_rel_bbox_area: float = 0.0
+    add_mask_channel: typing.Optional[bool] = None
     """
-    Minimum valid relative area (as a fraction of the image area) of a bounding box to keep it in the dataset for optimization.
+    Whether to add a mask channel to the image.
+    Change at your own risk. It is recommended to keep it at False unless you know what you are doing.
     """
 
 
-RunArgs = typing.Union[VisionRunArgs]
+RunArgs = typing.Union[ClassificationRunArgs, ObjectDetectionRunArgs]
 
 
 class AugmentationName(str, Enum):
@@ -176,13 +128,31 @@ class AugmentationName(str, Enum):
     GAUSSIAN_BLUR = "GaussianBlur"
 
 
-class Modality(str, Enum):
-    IMAGE = "Image"
-    RADAR = "Radar"
-    EKG = "EKG"
+class Domain(str, Enum):
+    RADAR = "RADAR"
+    VISION = "VISION"
+    SPEECH = "SPEECH"
+    TABULAR = "TABULAR"
+
+
+DOMAIN_TO_SUPPORTED_LABELING_TYPES = {
+    Domain.RADAR: [
+        LabelingType.SINGLE_LABEL_CLASSIFICATION,
+        LabelingType.OBJECT_DETECTION,
+    ],
+    Domain.VISION: [
+        LabelingType.SINGLE_LABEL_CLASSIFICATION,
+        LabelingType.OBJECT_DETECTION,
+        LabelingType.OBJECT_SEGMENTATION,
+        LabelingType.SEMANTIC_SEGMENTATION,
+        LabelingType.PANOPTIC_SEGMENTATION,
+    ],
+    Domain.SPEECH: [LabelingType.SPEECH_TO_TEXT],
+    Domain.TABULAR: [LabelingType.SINGLE_LABEL_CLASSIFICATION],
+}
 
 
-class OptimizationDataset(BaseModel):
+class QADataset(BaseModel):
     id: typing.Optional[int] = Field(default=None)
     """
     The ID of the dataset created on the server.
@@ -228,7 +198,7 @@ class OptimizationDataset(BaseModel):
     A full list of possible classes used in classification / object detection.
     It is currently required for clarity and performance.
     """
-    labeling_info: LabelingInfo
+    labeling_info: typing.Union[LabelingInfo, list[LabelingInfo]]
 
     augmentations: typing.Optional[list[AugmentationName]] = None
     """
@@ -236,21 +206,29 @@ class OptimizationDataset(BaseModel):
     For audio datasets, this field is ignored.
     If no value is provided, all augmentations are applied to vision datasets.
     """
-    modality: Modality = Modality.IMAGE
+    domain: Domain = Domain.VISION
     """
-    Used to define the modality of the dataset.
+    Used to define the domain of the dataset.
     Defaults to Image.
     """
 
     run_id: typing.Optional[str] = Field(default=None, init=False)
     """
-    The ID of the Dataset Optimization run created on the server.
+    The ID of the Dataset QA run created on the server.
     """
 
     status: typing.Optional[RunStatus] = None
 
     @model_validator(mode="after")
     def validate_dataset(self):
+        if self.domain not in DOMAIN_TO_SUPPORTED_LABELING_TYPES:
+            raise ValueError(
+                f"Domain {self.domain} is not supported. Supported domains are: {list(DOMAIN_TO_SUPPORTED_LABELING_TYPES.keys())}"
+            )
+        if self.labeling_type not in DOMAIN_TO_SUPPORTED_LABELING_TYPES[self.domain]:
+            raise ValueError(
+                f"Labeling type {self.labeling_type} is not supported for domain {self.domain}. Supported labeling types are: {DOMAIN_TO_SUPPORTED_LABELING_TYPES[self.domain]}"
+            )
         if self.storage_config is None and self.storage_config_id is None:
             raise ValueError(
                 "No dataset storage has been provided. Provide one via `storage_config` or `storage_config_id`"
@@ -267,65 +245,79 @@ class OptimizationDataset(BaseModel):
         ):
             raise ValueError("Language is only allowed for Speech-to-Text datasets.")
         if (
-            self.labeling_info.type == DatasetMetadataType.YOLO
+            not isinstance(self.labeling_info, list)
+            and self.labeling_info.type == DatasetMetadataType.YOLO
             and isinstance(self.labeling_info, YOLO)
             and (
                 self.labeling_info.data_yaml_url is not None
                 and self.classes is not None
             )
+        ) or (
+            isinstance(self.labeling_info, list)
+            and self.classes is not None
+            and any(
+                isinstance(info, YOLO) and info.data_yaml_url is not None
+                for info in self.labeling_info
+            )
         ):
             raise ValueError(
                 "Only one of `classes` or `labeling_info.data_yaml_url` should be provided for YOLO datasets"
             )
+        if self.storage_config:
+            validate_labeling_info(
+                self.labeling_type, self.labeling_info, self.storage_config
+            )
+        if self.data_root_url and self.storage_config:
+            validate_url(self.data_root_url, self.storage_config)
         return self
 
     @staticmethod
-    def get_by_id(dataset_id: int) -> "OptimizationDataset":
+    def get_by_id(dataset_id: int) -> "QADataset":
         """
-        Get a `OptimizationDataset` instance from the server by its ID
+        Get a `QADataset` instance from the server by its ID
 
         Args:
-            dataset_id: The ID of the `OptimizationDataset` instance to get
+            dataset_id: The ID of the `QADataset` instance to get
         """
         response = requests.get(
-            f"{API_HOST}/dataset-optimization/dataset/{dataset_id}",
+            f"{API_HOST}/dataset-qa/dataset/{dataset_id}",
             headers=get_headers(),
             timeout=READ_TIMEOUT,
         )
         raise_for_status_with_reason(response)
         dataset = response.json()
-        return OptimizationDataset(**dataset)
+        return QADataset(**dataset)
 
     @staticmethod
-    def get_by_name(name: str) -> "OptimizationDataset":
+    def get_by_name(name: str) -> "QADataset":
         """
-        Get a `OptimizationDataset` instance from the server by its name
+        Get a `QADataset` instance from the server by its name
 
         Args:
-            name: The name of the `OptimizationDataset` instance to get
+            name: The name of the `QADataset` instance to get
         """
         response = requests.get(
-            f"{API_HOST}/dataset-optimization/dataset/by-name/{name}",
+            f"{API_HOST}/dataset-qa/dataset/by-name/{name}",
             headers=get_headers(),
             timeout=READ_TIMEOUT,
         )
         raise_for_status_with_reason(response)
         dataset = response.json()
-        return OptimizationDataset(**dataset)
+        return QADataset(**dataset)
 
     @staticmethod
     def list_datasets(
         organization_id: typing.Optional[int] = None,
-    ) -> list["DataOptimizationDatasetOut"]:
+    ) -> list["QADatasetOut"]:
         """
-        Lists all the optimization datasets created by user's default organization
+        Lists all the datasets created by user's default organization
         or the `organization_id` passed
 
         Args:
             organization_id: The ID of the organization to list the datasets for.
         """
         response = requests.get(
-            f"{API_HOST}/dataset-optimization/dataset/",
+            f"{API_HOST}/dataset-qa/dataset/",
             params={"dataset_organization_id": organization_id},
             headers=get_headers(),
             timeout=READ_TIMEOUT,
@@ -333,7 +325,7 @@ class OptimizationDataset(BaseModel):
         raise_for_status_with_reason(response)
         datasets = response.json()
         return [
-            DataOptimizationDatasetOut(
+            QADatasetOut(
                 **ds,
             )
             for ds in datasets
@@ -342,17 +334,17 @@ class OptimizationDataset(BaseModel):
     @staticmethod
     def list_runs(
         organization_id: typing.Optional[int] = None,
-    ) -> list["DataOptimizationRunOut"]:
+    ) -> list["DataQARunOut"]:
         """
-        Lists all the `OptimizationDataset` instances created by user's default organization
+        Lists all the `QADataset` instances created by user's default organization
         or the `organization_id` passed
-        Note: The return type is `list[dict]` and not `list[OptimizationDataset]`
+        Note: The return type is `list[dict]` and not `list[QADataset]`
 
         Args:
             organization_id: The ID of the organization to list the datasets for.
         """
         response = requests.get(
-            f"{API_HOST}/dataset-optimization/run/list",
+            f"{API_HOST}/dataset-qa/run/list",
             params={"dataset_organization_id": organization_id},
             headers=get_headers(),
             timeout=READ_TIMEOUT,
@@ -360,7 +352,7 @@ class OptimizationDataset(BaseModel):
         raise_for_status_with_reason(response)
         runs = response.json()
         return [
-            DataOptimizationRunOut(
+            DataQARunOut(
                 **run,
             )
             for run in runs
@@ -369,13 +361,13 @@ class OptimizationDataset(BaseModel):
     @staticmethod
     def delete_by_id(dataset_id: int) -> None:
         """
-        Deletes a `OptimizationDataset` instance from the server by its ID
+        Deletes a `QADataset` instance from the server by its ID
 
         Args:
-            dataset_id: The ID of the `OptimizationDataset` instance to delete
+            dataset_id: The ID of the `QADataset` instance to delete
         """
         response = requests.delete(
-            f"{API_HOST}/dataset-optimization/dataset/{dataset_id}",
+            f"{API_HOST}/dataset-qa/dataset/{dataset_id}",
             headers=get_headers(),
             timeout=MODIFY_TIMEOUT,
         )
@@ -384,14 +376,14 @@ class OptimizationDataset(BaseModel):
 
     def delete(self, storage_config=True) -> None:
         """
-        Deletes the active `OptimizationDataset` instance from the server.
-        It can only be used on a `OptimizationDataset` instance that has been created.
+        Deletes the active `QADataset` instance from the server.
+        It can only be used on a `QADataset` instance that has been created.
 
         Args:
-            storage_config: If True, the `OptimizationDataset`'s `StorageConfig` will also be deleted
+            storage_config: If True, the `QADataset`'s `StorageConfig` will also be deleted
 
         Note: If `storage_config` is not set to `False` then the `storage_config_id` must be set
-        This can either be set manually or by creating the `StorageConfig` instance via the `OptimizationDataset`'s
+        This can either be set manually or by creating the `StorageConfig` instance via the `QADataset`'s
         `create` method
         """
         if storage_config:
@@ -408,7 +400,7 @@ class OptimizationDataset(BaseModel):
         replace_if_exists: bool = False,
     ) -> int:
         """
-        Create a `OptimizationDataset` instance on the server.
+        Create a `QADataset` instance on the server.
         If the `storage_config_id` field is not set, the storage config will also be created and the field will be set.
 
         Args:
@@ -417,7 +409,7 @@ class OptimizationDataset(BaseModel):
                 (this is determined by a dataset of the same name in the same organization).
 
         Returns:
-            The ID of the created `OptimizationDataset` instance
+            The ID of the created `QADataset` instance
         """
         if self.storage_config is None and self.storage_config_id is None:
             raise ValueError("No dataset storage has been provided")
@@ -442,7 +434,7 @@ class OptimizationDataset(BaseModel):
         model_dict = self.model_dump(mode="json")
         # ⬆️ Get dict of model fields from Pydantic model instance
         dataset_response = requests.post(
-            f"{API_HOST}/dataset-optimization/dataset/",
+            f"{API_HOST}/dataset-qa/dataset/",
             json={
                 **{k: model_dict[k] for k in model_dict.keys() - {"storage_config"}},
                 "organization_id": organization_id,
@@ -459,17 +451,17 @@ class OptimizationDataset(BaseModel):
         return self.id
 
     @staticmethod
-    def launch_optimization_run(
+    def launch_qa_run(
         dataset_id: int,
         organization_id: typing.Optional[int] = None,
         run_args: typing.Optional[RunArgs] = None,
     ) -> str:
         """
-        Run the dataset optimization process on the server using the dataset with the given ID
+        Run the dataset QA process on the server using the dataset with the given ID
         i.e. `dataset_id`.
 
         Args:
-            dataset_id: The ID of the dataset to run optimization on.
+            dataset_id: The ID of the dataset to run QA on.
 
         Returns:
             ID of the run (`run_id`).
@@ -480,7 +472,7 @@ class OptimizationDataset(BaseModel):
         if run_args:
             run_info["run_args"] = run_args.model_dump(mode="json")
         run_response = requests.post(
-            f"{API_HOST}/dataset-optimization/run/{dataset_id}",
+            f"{API_HOST}/dataset-qa/run/{dataset_id}",
             json=run_info if len(run_info) > 0 else None,
             headers=get_headers(),
             timeout=MODIFY_TIMEOUT,
@@ -491,12 +483,16 @@ class OptimizationDataset(BaseModel):
     def _validate_run_args(self, run_args: RunArgs) -> None:
         if self.labeling_type == LabelingType.SPEECH_TO_TEXT:
             raise Exception("Speech to text cannot have `run_args` set")
-        if self.labeling_type != LabelingType.OBJECT_DETECTION and any(
-            (
-                run_args.min_abs_bbox_size != 0,
-                run_args.min_abs_bbox_area != 0,
-                run_args.min_rel_bbox_size != 0,
-                run_args.min_rel_bbox_area != 0,
+        if (
+            self.labeling_type != LabelingType.OBJECT_DETECTION
+            and isinstance(run_args, ObjectDetectionRunArgs)
+            and any(
+                (
+                    run_args.min_abs_bbox_size != 0,
+                    run_args.min_abs_bbox_area != 0,
+                    run_args.min_rel_bbox_size != 0,
+                    run_args.min_rel_bbox_area != 0,
+                )
             )
         ):
             raise Exception(
@@ -505,7 +501,7 @@ class OptimizationDataset(BaseModel):
                 + f"labeling type {self.labeling_type}"
             )
 
-    def run_optimization(
+    def run_qa(
         self,
         organization_id: typing.Optional[int] = None,
         replace_dataset_if_exists: bool = False,
@@ -513,13 +509,13 @@ class OptimizationDataset(BaseModel):
     ) -> str:
         """
         If the dataset was not created on the server yet, it is created.
-        Run the dataset optimization process on the server using the active `OptimizationDataset` instance
+        Run the dataset QA process on the server using the active `QADataset` instance
 
         Args:
-            organization_id: The ID of the organization to run the optimization for.
+            organization_id: The ID of the organization to run the QA for.
             replace_dataset_if_exists: If True, the dataset will be replaced if it already exists
                 (this is determined by a dataset of the same name in the same organization).
-            run_args: The run arguments to use for the optimization run
+            run_args: The run arguments to use for the QA run
 
         Returns:
             An ID of the run (`run_id`) and stores that `run_id` on the instance
@@ -529,7 +525,7 @@ class OptimizationDataset(BaseModel):
                 self.id = self.create(replace_if_exists=replace_dataset_if_exists)
             if run_args is not None:
                 self._validate_run_args(run_args)
-            run_id = self.launch_optimization_run(self.id, organization_id, run_args)
+            run_id = self.launch_qa_run(self.id, organization_id, run_args)
             self.run_id = run_id
             logger.info("Started the run with ID: %s", run_id)
             return run_id
@@ -567,7 +563,7 @@ class OptimizationDataset(BaseModel):
             for sse in iter_sse_retrying(
                 client,
                 "GET",
-                f"{API_HOST}/dataset-optimization/run/{run_id}",
+                f"{API_HOST}/dataset-qa/run/{run_id}",
                 headers=get_headers(),
             ):
                 if sse.event == "ping":
@@ -593,39 +589,46 @@ class OptimizationDataset(BaseModel):
                         raise HirundoError("Unknown error")
                 yield data
         if not last_event or last_event["data"]["state"] == RunStatus.PENDING.value:
-            OptimizationDataset._check_run_by_id(run_id, retry + 1)
+            QADataset._check_run_by_id(run_id, retry + 1)
+
+    @staticmethod
+    def _handle_failure(iteration: dict):
+        if iteration["result"]:
+            raise HirundoError(f"QA run failed with error: {iteration['result']}")
+        else:
+            raise HirundoError("QA run failed with an unknown error in _handle_failure")
 
     @staticmethod
     @overload
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: typing.Literal[True]
-    ) -> typing.Optional[DatasetOptimizationResults]: ...
+    ) -> typing.Optional[DatasetQAResults]: ...
 
     @staticmethod
     @overload
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: typing.Literal[False] = False
-    ) -> DatasetOptimizationResults: ...
+    ) -> DatasetQAResults: ...
 
     @staticmethod
     @overload
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: bool
-    ) -> typing.Optional[DatasetOptimizationResults]: ...
+    ) -> typing.Optional[DatasetQAResults]: ...
 
     @staticmethod
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: bool = False
-    ) -> typing.Optional[DatasetOptimizationResults]:
+    ) -> typing.Optional[DatasetQAResults]:
         """
         Check the status of a run given its ID
 
         Args:
-            run_id: The `run_id` produced by a `run_optimization` call
+            run_id: The `run_id` produced by a `run_qa` call
             stop_on_manual_approval: If True, the function will return `None` if the run is awaiting manual approval
 
         Returns:
-            A DatasetOptimizationResults object with the results of the optimization run
+            A DatasetQAResults object with the results of the QA run
 
         Raises:
             HirundoError: If the maximum number of retries is reached or if the run fails
@@ -633,7 +636,7 @@ class OptimizationDataset(BaseModel):
         logger.debug("Checking run with ID: %s", run_id)
         with logging_redirect_tqdm():
             t = tqdm(total=100.0)
-            for iteration in OptimizationDataset._check_run_by_id(run_id):
+            for iteration in QADataset._check_run_by_id(run_id):
                 if iteration["state"] in STATUS_TO_PROGRESS_MAP:
                     t.set_description(STATUS_TO_TEXT_MAP[iteration["state"]])
                     t.n = STATUS_TO_PROGRESS_MAP[iteration["state"]]
@@ -644,13 +647,15 @@ class OptimizationDataset(BaseModel):
                         RunStatus.REJECTED.value,
                         RunStatus.REVOKED.value,
                     ]:
-                        raise HirundoError(
-                            f"Optimization run failed with error: {iteration['result']}"
+                        logger.error(
+                            "State is failure, rejected, or revoked: %s",
+                            iteration["state"],
                         )
+                        QADataset._handle_failure(iteration)
                     elif iteration["state"] == RunStatus.SUCCESS.value:
                         t.close()
                         zip_temporary_url = iteration["result"]
-                        logger.debug("Optimization run completed. Downloading results")
+                        logger.debug("QA run completed. Downloading results")
 
                         return download_and_extract_zip(
                             run_id,
@@ -682,7 +687,7 @@ class OptimizationDataset(BaseModel):
                             stage = "Unknown progress state"
                             current_progress_percentage = t.n  # Keep the same progress
                         desc = (
-                            "Optimization run completed. Uploading results"
+                            "QA run completed. Uploading results"
                             if current_progress_percentage == 100.0
                             else stage
                         )
@@ -690,26 +695,26 @@ class OptimizationDataset(BaseModel):
                         t.n = current_progress_percentage
                         logger.debug("Setting progress to %s", t.n)
                         t.refresh()
-        raise HirundoError("Optimization run failed with an unknown error")
+        raise HirundoError("QA run failed with an unknown error in check_run_by_id")
 
     @overload
     def check_run(
         self, stop_on_manual_approval: typing.Literal[True]
-    ) -> typing.Optional[DatasetOptimizationResults]: ...
+    ) -> typing.Optional[DatasetQAResults]: ...
 
     @overload
     def check_run(
         self, stop_on_manual_approval: typing.Literal[False] = False
-    ) -> DatasetOptimizationResults: ...
+    ) -> DatasetQAResults: ...
 
     def check_run(
         self, stop_on_manual_approval: bool = False
-    ) -> typing.Optional[DatasetOptimizationResults]:
+    ) -> typing.Optional[DatasetQAResults]:
         """
         Check the status of the current active instance's run.
 
         Returns:
-            A pandas DataFrame with the results of the optimization run
+            A pandas DataFrame with the results of the QA run
 
         """
         if not self.run_id:
@@ -726,7 +731,7 @@ class OptimizationDataset(BaseModel):
         This generator will produce values to show progress of the run.
 
         Args:
-            run_id: The `run_id` produced by a `run_optimization` call
+            run_id: The `run_id` produced by a `run_qa` call
             retry: A number used to track the number of retries to limit re-checks. *Do not* provide this value manually.
 
         Yields:
@@ -745,7 +750,7 @@ class OptimizationDataset(BaseModel):
             async_iterator = await aiter_sse_retrying(
                 client,
                 "GET",
-                f"{API_HOST}/dataset-optimization/run/{run_id}",
+                f"{API_HOST}/dataset-qa/run/{run_id}",
                 headers=get_headers(),
             )
             async for sse in async_iterator:
@@ -761,7 +766,7 @@ class OptimizationDataset(BaseModel):
                 last_event = json.loads(sse.data)
                 yield last_event["data"]
         if not last_event or last_event["data"]["state"] == RunStatus.PENDING.value:
-            OptimizationDataset.acheck_run_by_id(run_id, retry + 1)
+            QADataset.acheck_run_by_id(run_id, retry + 1)
 
     async def acheck_run(self) -> AsyncGenerator[dict, None]:
         """
@@ -785,16 +790,14 @@ class OptimizationDataset(BaseModel):
     @staticmethod
     def cancel_by_id(run_id: str) -> None:
         """
-        Cancel the dataset optimization run for the given `run_id`.
+        Cancel the dataset QA run for the given `run_id`.
 
         Args:
             run_id: The ID of the run to cancel
         """
-        if not run_id:
-            raise ValueError("No run has been started")
         logger.info("Cancelling run with ID: %s", run_id)
         response = requests.delete(
-            f"{API_HOST}/dataset-optimization/run/{run_id}",
+            f"{API_HOST}/dataset-qa/run/{run_id}",
             headers=get_headers(),
             timeout=MODIFY_TIMEOUT,
         )
@@ -808,8 +811,32 @@ class OptimizationDataset(BaseModel):
             raise ValueError("No run has been started")
         self.cancel_by_id(self.run_id)
 
+    @staticmethod
+    def archive_run_by_id(run_id: str) -> None:
+        """
+        Archive the dataset QA run for the given `run_id`.
+
+        Args:
+            run_id: The ID of the run to archive
+        """
+        logger.info("Archiving run with ID: %s", run_id)
+        response = requests.patch(
+            f"{API_HOST}/dataset-qa/run/archive/{run_id}",
+            headers=get_headers(),
+            timeout=MODIFY_TIMEOUT,
+        )
+        raise_for_status_with_reason(response)
+
+    def archive(self) -> None:
+        """
+        Archive the current active instance's run.
+        """
+        if not self.run_id:
+            raise ValueError("No run has been started")
+        self.archive_run_by_id(self.run_id)
+
 
-class DataOptimizationDatasetOut(BaseModel):
+class QADatasetOut(BaseModel):
     id: int
 
     name: str
@@ -820,7 +847,7 @@ class DataOptimizationDatasetOut(BaseModel):
     data_root_url: HirundoUrl
 
     classes: typing.Optional[list[str]] = None
-    labeling_info: LabelingInfo
+    labeling_info: typing.Union[LabelingInfo, list[LabelingInfo]]
 
     organization_id: typing.Optional[int]
     creator_id: typing.Optional[int]
@@ -828,7 +855,7 @@ class DataOptimizationDatasetOut(BaseModel):
     updated_at: datetime.datetime
 
 
-class DataOptimizationRunOut(BaseModel):
+class DataQARunOut(BaseModel):
     id: int
     name: str
     dataset_id: int